diff --git a/docs/changes/newsfragments/7725.improved b/docs/changes/newsfragments/7725.improved new file mode 100644 index 00000000000..877254dca67 --- /dev/null +++ b/docs/changes/newsfragments/7725.improved @@ -0,0 +1,6 @@ +Parameters using ``has_control_of`` are now correctly handled when exporting to +xarray. Controlled parameters are no longer treated as independent top-level +parameters, preventing duplicate data rows. Additionally, inferred parameters +are now included as data variables in the xarray dataset when exporting via the +pandas-based path, and a warning is logged when the inferred parameter data size +does not match the expected xarray dataset dimensions. diff --git a/src/qcodes/dataset/descriptions/dependencies.py b/src/qcodes/dataset/descriptions/dependencies.py index 3ec82477dc0..43c374c5703 100644 --- a/src/qcodes/dataset/descriptions/dependencies.py +++ b/src/qcodes/dataset/descriptions/dependencies.py @@ -287,6 +287,16 @@ def top_level_parameters(self) -> tuple[ParamSpecBase, ...]: for node_id, in_degree in self._dependency_subgraph.in_degree if in_degree == 0 } + # Parameters that are inferred from other parameters (have outgoing + # edges in the inference subgraph) should not be independent top-level + # parameters, since their data is part of the tree of the parameter + # they are inferred from. + parameters_inferred_from_others = { + self._node_to_paramspec(node_id) + for node_id, out_degree in self._inference_subgraph.out_degree + if out_degree > 0 + } + dependency_top_level = dependency_top_level - parameters_inferred_from_others standalone_top_level = { self._node_to_paramspec(node_id) for node_id, degree in self._graph.degree diff --git a/src/qcodes/dataset/exporters/export_to_xarray.py b/src/qcodes/dataset/exporters/export_to_xarray.py index 1582d15dd03..1734831dd82 100644 --- a/src/qcodes/dataset/exporters/export_to_xarray.py +++ b/src/qcodes/dataset/exporters/export_to_xarray.py @@ -6,6 +6,7 @@ from math import prod from typing import TYPE_CHECKING, Literal +import numpy as np from packaging import version as p_version from qcodes.dataset.linked_datasets.links import links_to_str @@ -61,6 +62,68 @@ def _calculate_index_shape(idx: pd.Index | pd.MultiIndex) -> dict[Hashable, int] return expanded_shape +def _add_inferred_data_vars( + dataset: DataSetProtocol, + name: str, + sub_dict: Mapping[str, npt.NDArray], + xr_dataset: xr.Dataset, +) -> xr.Dataset: + """Add inferred parameters as data variables to an xarray dataset. + + Parameters that are inferred from the top-level measurement parameter + and present in sub_dict but not yet in the dataset are added as data + variables along the existing dimensions. + """ + + interdeps = dataset.description.interdeps + meas_paramspec = interdeps.graph.nodes[name]["value"] + _, deps, inferred = interdeps.all_parameters_in_tree_by_group(meas_paramspec) + + dep_names = {dep.name for dep in deps} + dims = tuple(d for d in xr_dataset.dims) + + for inf in inferred: + if inf.name in dep_names: + continue + if inf.name in xr_dataset: + continue + if inf.name not in sub_dict: + continue + + inf_data = sub_dict[inf.name] + if inf_data.dtype == np.dtype("O"): + try: + flat = np.concatenate(inf_data) + except ValueError: + flat = inf_data.ravel() + else: + flat = inf_data.ravel() + + # Only add if the flattened data can be reshaped to the dataset + # dimensions. This is more robust than checking individual parent + # sizes because an inferred parameter may have multiple parents + # with different sizes. + expected_shape = tuple(xr_dataset.sizes[d] for d in dims) + expected_size = prod(expected_shape) + if flat.shape[0] == expected_size: + xr_dataset[inf.name] = (dims, flat.reshape(expected_shape)) + else: + _LOG.warning( + "Cannot add inferred parameter '%s' to xarray dataset for '%s' " + "(run_id=%s): data size %d does not match the dataset " + "dimensions %s (size %d). This is likely a user error in the " + "measurement setup.", + inf.name, + name, + dataset.run_id, + flat.shape[0], + dict(zip(dims, expected_shape)), + expected_size, + ) + + return xr_dataset + + def _load_to_xarray_dataset_dict_no_metadata( dataset: DataSetProtocol, datadict: Mapping[str, Mapping[str, npt.NDArray]], @@ -100,7 +163,9 @@ def _load_to_xarray_dataset_dict_no_metadata( interdeps=dataset.description.interdeps, dependent_parameter=name, ).to_xarray() - xr_dataset_dict[name] = xr_dataset + xr_dataset_dict[name] = _add_inferred_data_vars( + dataset, name, sub_dict, xr_dataset + ) elif index_is_unique: df = _data_to_dataframe( sub_dict, @@ -108,9 +173,12 @@ def _load_to_xarray_dataset_dict_no_metadata( interdeps=dataset.description.interdeps, dependent_parameter=name, ) - xr_dataset_dict[name] = _xarray_data_set_from_pandas_multi_index( + xr_dataset = _xarray_data_set_from_pandas_multi_index( dataset, use_multi_index, name, df, index ) + xr_dataset_dict[name] = _add_inferred_data_vars( + dataset, name, sub_dict, xr_dataset + ) else: df = _data_to_dataframe( sub_dict, @@ -118,7 +186,10 @@ def _load_to_xarray_dataset_dict_no_metadata( interdeps=dataset.description.interdeps, dependent_parameter=name, ) - xr_dataset_dict[name] = df.reset_index().to_xarray() + xr_dataset = df.reset_index().to_xarray() + xr_dataset_dict[name] = _add_inferred_data_vars( + dataset, name, sub_dict, xr_dataset + ) return xr_dataset_dict diff --git a/tests/dataset/test_inferred_multiple_parents.py b/tests/dataset/test_inferred_multiple_parents.py new file mode 100644 index 00000000000..770044ab354 --- /dev/null +++ b/tests/dataset/test_inferred_multiple_parents.py @@ -0,0 +1,374 @@ +"""Tests for _add_inferred_data_vars with multiple inferred-from parents. + +These tests verify that the inferred parameter's data size is checked +against the xr_dataset dimensions rather than individual parent sizes. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING +from unittest.mock import MagicMock + +import numpy as np +import numpy.testing as npt +import xarray as xr + +from qcodes.dataset.descriptions.dependencies import InterDependencies_ +from qcodes.dataset.exporters.export_to_xarray import _add_inferred_data_vars +from qcodes.parameters import ParamSpecBase + +if TYPE_CHECKING: + import pytest + + +def _make_mock_dataset( + interdeps: InterDependencies_, + run_id: int = 1, +) -> MagicMock: + """Create a minimal mock DataSetProtocol with the given interdeps.""" + ds = MagicMock() + ds.description.interdeps = interdeps + ds.run_id = run_id + return ds + + +def _make_interdeps( + *, + deps: dict[ParamSpecBase, tuple[ParamSpecBase, ...]], + inferences: dict[ParamSpecBase, tuple[ParamSpecBase, ...]], +) -> InterDependencies_: + return InterDependencies_(dependencies=deps, inferences=inferences) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has ONE parent, sizes match → included +# (baseline sanity check) +# --------------------------------------------------------------------------- +class TestSingleParentBaseline: + def test_single_parent_matching_size_is_included(self) -> None: + """An inferred param whose data matches its single parent is added.""" + sp = ParamSpecBase("sp", "numeric") + meas = ParamSpecBase("meas", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={meas: (sp,)}, + inferences={inf: (meas,)}, + ) + ds = _make_mock_dataset(interdeps) + + n = 10 + sub_dict: dict[str, np.ndarray] = { + "sp": np.arange(n, dtype=float), + "meas": np.random.default_rng().standard_normal(n), + "inf_param": np.linspace(0, 1, n), + } + + xr_ds = xr.Dataset( + {"meas": (("sp",), sub_dict["meas"])}, + coords={"sp": sub_dict["sp"]}, + ) + + result = _add_inferred_data_vars(ds, "meas", sub_dict, xr_ds) + + assert "inf_param" in result.data_vars + npt.assert_array_almost_equal(result["inf_param"].values, sub_dict["inf_param"]) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has TWO parents, data matches BOTH +# → should always be included regardless of strategy +# --------------------------------------------------------------------------- +class TestMultipleParentsAllMatch: + def test_inferred_matches_all_parents_is_included(self) -> None: + """When data size matches all parents, the inferred param is added.""" + sp = ParamSpecBase("sp", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp,), parent2: (sp,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n = 10 + sub_dict: dict[str, np.ndarray] = { + "sp": np.arange(n, dtype=float), + "parent1": np.random.default_rng().standard_normal(n), + "parent2": np.random.default_rng().standard_normal(n), + "inf_param": np.linspace(0, 1, n), + } + + xr_ds = xr.Dataset( + { + "parent1": (("sp",), sub_dict["parent1"]), + "parent2": (("sp",), sub_dict["parent2"]), + }, + coords={"sp": sub_dict["sp"]}, + ) + + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + assert "inf_param" in result.data_vars + npt.assert_array_almost_equal(result["inf_param"].values, sub_dict["inf_param"]) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has TWO parents with DIFFERENT sizes, +# data matches only the FIRST parent +# → current "match any" includes it; "match all" would reject it +# --------------------------------------------------------------------------- +class TestMultipleParentsOnlyFirstMatches: + def test_inferred_matches_first_parent_only(self) -> None: + """Data matches parent1 (size 10) but not parent2 (size 5). + + Current behavior: included (matches any parent). + If "match all" were required, this would NOT be included. + """ + sp1 = ParamSpecBase("sp1", "numeric") + sp2 = ParamSpecBase("sp2", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp1,), parent2: (sp2,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n1, n2 = 10, 5 + sub_dict: dict[str, np.ndarray] = { + "sp1": np.arange(n1, dtype=float), + "sp2": np.arange(n2, dtype=float), + "parent1": np.random.default_rng().standard_normal(n1), + "parent2": np.random.default_rng().standard_normal(n2), + # inf_param has the same size as parent1 + "inf_param": np.linspace(0, 1, n1), + } + + xr_ds = xr.Dataset( + {"parent1": (("sp1",), sub_dict["parent1"])}, + coords={"sp1": sub_dict["sp1"]}, + ) + + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + # Current behavior: included because it matches parent1 + assert "inf_param" in result.data_vars + npt.assert_array_almost_equal(result["inf_param"].values, sub_dict["inf_param"]) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has TWO parents with DIFFERENT sizes, +# data matches only the SECOND parent +# → current "match any" includes it; "match all" would reject it +# --------------------------------------------------------------------------- +class TestMultipleParentsOnlySecondMatches: + def test_inferred_matches_second_parent_not_dataset_dims_warns( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Data matches parent2 (size 5) but not the dataset dims (size 10). + + The inferred param should NOT be included because its data cannot + be reshaped to the xr_dataset dimensions. A warning is emitted. + """ + sp1 = ParamSpecBase("sp1", "numeric") + sp2 = ParamSpecBase("sp2", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp1,), parent2: (sp2,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n1, n2 = 10, 5 + sub_dict: dict[str, np.ndarray] = { + "sp1": np.arange(n1, dtype=float), + "sp2": np.arange(n2, dtype=float), + "parent1": np.random.default_rng().standard_normal(n1), + "parent2": np.random.default_rng().standard_normal(n2), + # inf_param has the same size as parent2 but NOT the dataset dims + "inf_param": np.linspace(0, 1, n2), + } + + xr_ds = xr.Dataset( + {"parent1": (("sp1",), sub_dict["parent1"])}, + coords={"sp1": sub_dict["sp1"]}, + ) + + with caplog.at_level( + logging.WARNING, logger="qcodes.dataset.exporters.export_to_xarray" + ): + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + assert "inf_param" not in result.data_vars + assert any( + "Cannot add inferred parameter 'inf_param'" in msg + for msg in caplog.messages + ) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has TWO parents, data matches NEITHER +# → should be excluded and emit a warning +# --------------------------------------------------------------------------- +class TestMultipleParentsNoneMatch: + def test_inferred_matches_no_parent_warns( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Data size doesn't match any parent → warning, not included.""" + sp = ParamSpecBase("sp", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp,), parent2: (sp,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n = 10 + sub_dict: dict[str, np.ndarray] = { + "sp": np.arange(n, dtype=float), + "parent1": np.random.default_rng().standard_normal(n), + "parent2": np.random.default_rng().standard_normal(n), + # inf_param has a completely different size + "inf_param": np.linspace(0, 1, 7), + } + + xr_ds = xr.Dataset( + {"parent1": (("sp",), sub_dict["parent1"])}, + coords={"sp": sub_dict["sp"]}, + ) + + with caplog.at_level( + logging.WARNING, logger="qcodes.dataset.exporters.export_to_xarray" + ): + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + assert "inf_param" not in result.data_vars + assert any( + "Cannot add inferred parameter 'inf_param'" in msg + for msg in caplog.messages + ) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has TWO parents, only one is in sub_dict +# → should match against the available parent +# --------------------------------------------------------------------------- +class TestMultipleParentsOneUnavailable: + def test_matches_available_parent_ignores_missing(self) -> None: + """When one parent is not in sub_dict, the other is still checked.""" + sp = ParamSpecBase("sp", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp,), parent2: (sp,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n = 10 + sub_dict: dict[str, np.ndarray] = { + "sp": np.arange(n, dtype=float), + "parent1": np.random.default_rng().standard_normal(n), + # parent2 is NOT in sub_dict + "inf_param": np.linspace(0, 1, n), + } + + xr_ds = xr.Dataset( + {"parent1": (("sp",), sub_dict["parent1"])}, + coords={"sp": sub_dict["sp"]}, + ) + + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + assert "inf_param" in result.data_vars + npt.assert_array_almost_equal(result["inf_param"].values, sub_dict["inf_param"]) + + def test_warns_when_only_available_parent_mismatches( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """One parent missing, the other has wrong size → warning.""" + sp = ParamSpecBase("sp", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp,), parent2: (sp,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n = 10 + sub_dict: dict[str, np.ndarray] = { + "sp": np.arange(n, dtype=float), + "parent1": np.random.default_rng().standard_normal(n), + # parent2 missing, inf_param has wrong size + "inf_param": np.linspace(0, 1, 7), + } + + xr_ds = xr.Dataset( + {"parent1": (("sp",), sub_dict["parent1"])}, + coords={"sp": sub_dict["sp"]}, + ) + + with caplog.at_level( + logging.WARNING, logger="qcodes.dataset.exporters.export_to_xarray" + ): + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + assert "inf_param" not in result.data_vars + assert any( + "Cannot add inferred parameter 'inf_param'" in msg + for msg in caplog.messages + ) + + +# --------------------------------------------------------------------------- +# Scenario: inferred param has TWO parents of SAME size, both match +# → should be included; the "match any" and "match all" give same result +# --------------------------------------------------------------------------- +class TestMultipleParentsSameSizeAllMatch: + def test_both_parents_same_size_included(self) -> None: + """Both parents have same size and match → included either way.""" + sp = ParamSpecBase("sp", "numeric") + parent1 = ParamSpecBase("parent1", "numeric") + parent2 = ParamSpecBase("parent2", "numeric") + inf = ParamSpecBase("inf_param", "numeric") + + interdeps = _make_interdeps( + deps={parent1: (sp,), parent2: (sp,)}, + inferences={inf: (parent1, parent2)}, + ) + ds = _make_mock_dataset(interdeps) + + n = 10 + sub_dict: dict[str, np.ndarray] = { + "sp": np.arange(n, dtype=float), + "parent1": np.random.default_rng().standard_normal(n), + "parent2": np.random.default_rng().standard_normal(n), + "inf_param": np.linspace(0, 1, n), + } + + xr_ds = xr.Dataset( + {"parent1": (("sp",), sub_dict["parent1"])}, + coords={"sp": sub_dict["sp"]}, + ) + + result = _add_inferred_data_vars(ds, "parent1", sub_dict, xr_ds) + + assert "inf_param" in result.data_vars + npt.assert_array_almost_equal(result["inf_param"].values, sub_dict["inf_param"]) diff --git a/tests/dataset/test_parameter_with_setpoints_has_control.py b/tests/dataset/test_parameter_with_setpoints_has_control.py new file mode 100644 index 00000000000..08e641a89bd --- /dev/null +++ b/tests/dataset/test_parameter_with_setpoints_has_control.py @@ -0,0 +1,196 @@ +import logging +from typing import TYPE_CHECKING + +import numpy as np +import numpy.testing as npt +import xarray as xr + +from qcodes.dataset import Measurement +from qcodes.dataset.exporters.export_to_xarray import _add_inferred_data_vars +from qcodes.parameters import ManualParameter, Parameter, ParameterWithSetpoints +from qcodes.validators import Arrays + +if TYPE_CHECKING: + import pytest + + from qcodes.dataset.experiment_container import Experiment + + +def _make_controlled_setpoints( + name: str, + controlled: Parameter, + **kwargs: object, +) -> ParameterWithSetpoints: + """Create a ParameterWithSetpoints that infers ``controlled`` via unpack_self.""" + + class _ControlledSetpoints(ParameterWithSetpoints): + def unpack_self(self, value): # type: ignore[override] + res = super().unpack_self(value) + res.append((controlled, controlled())) + return res + + p = _ControlledSetpoints(name, **kwargs) # type: ignore[arg-type] + p.has_control_of.add(controlled) + return p + + +def test_parameter_with_setpoints_has_control(experiment: "Experiment"): + mp_data = np.arange(10) + p1_data = np.linspace(-1, 1, 10) + + mp = ManualParameter("mp", vals=Arrays(shape=(10,)), initial_value=mp_data) + p1 = ParameterWithSetpoints( + "p1", vals=Arrays(shape=(10,)), setpoints=(mp,), set_cmd=None + ) + p2 = _make_controlled_setpoints( + "p2", p1, vals=Arrays(shape=(10,)), setpoints=(mp,), set_cmd=None + ) + + p1(p1_data) + p2_data = np.random.default_rng().standard_normal(10) + p2(p2_data) + + meas = Measurement() + meas.register_parameter(p2) + + # Only p2 should be top-level; p1 is inferred from p2 + interdeps = meas._interdeps + top_level_names = [p.name for p in interdeps.top_level_parameters] + assert top_level_names == ["p2"] + + with meas.run() as ds: + ds.add_result((p2, p2())) + + # Verify raw parameter data has exactly one row per parameter + raw_data = ds.dataset.get_parameter_data() + assert list(raw_data.keys()) == ["p2"], "Only p2 should be a top-level result" + for name, arr in raw_data["p2"].items(): + assert arr.shape == (1, 10), ( + f"Expected shape (1, 10) for {name}, got {arr.shape}" + ) + + xds = ds.dataset.to_xarray_dataset() + + # mp should be the only dimension (not a generic 'index') + assert list(xds.sizes.keys()) == ["mp"] + assert xds.sizes["mp"] == 10 + + # mp values used as coordinate axis + npt.assert_array_equal(xds.coords["mp"].values, mp_data) + + # p2 is the primary data variable with correct values + assert "p2" in xds.data_vars + npt.assert_array_almost_equal(xds["p2"].values, p2_data) + + # p1 is included as a data variable (inferred from p2) with correct values + assert "p1" in xds.data_vars + npt.assert_array_almost_equal(xds["p1"].values, p1_data) + + # p1 data is also retrievable from the raw parameter data + npt.assert_array_almost_equal(raw_data["p2"]["p1"].ravel(), p1_data) + + +def test_parameter_with_setpoints_has_control_2d(experiment: "Experiment"): + """Test that an inferred parameter with the same size as its parent + but different from the full dimension product is correctly included.""" + + n_x = 3 + n_y = 4 + mp_x_data = np.arange(n_x, dtype=float) + mp_y_data = np.arange(n_y, dtype=float) + + mp_x = ManualParameter("mp_x", initial_value=0.0) + mp_y = ManualParameter("mp_y", vals=Arrays(shape=(n_y,)), initial_value=mp_y_data) + + p1 = ParameterWithSetpoints( + "p1", vals=Arrays(shape=(n_y,)), setpoints=(mp_y,), set_cmd=None + ) + p2 = _make_controlled_setpoints( + "p2", p1, vals=Arrays(shape=(n_y,)), setpoints=(mp_y,), set_cmd=None + ) + + meas = Measurement() + meas.register_parameter(p2, setpoints=(mp_x,)) + + p1_all = [] + p2_all = [] + + with meas.run() as ds: + for x_val in mp_x_data: + mp_x(x_val) + p1_row = np.linspace(-1, 1, n_y) + x_val + p1(p1_row) + p2_row = np.random.default_rng().standard_normal(n_y) + p2(p2_row) + p1_all.append(p1_row) + p2_all.append(p2_row) + ds.add_result((mp_x, mp_x()), (p2, p2())) + + p1_all_arr = np.array(p1_all) + p2_all_arr = np.array(p2_all) + + xds = ds.dataset.to_xarray_dataset() + + # Should have 2 dimensions: mp_x and mp_y + assert set(xds.sizes.keys()) == {"mp_x", "mp_y"} + assert xds.sizes["mp_x"] == n_x + assert xds.sizes["mp_y"] == n_y + + # p2 is the primary data variable + assert "p2" in xds.data_vars + npt.assert_array_almost_equal(xds["p2"].values, p2_all_arr) + + # p1 is included as a data variable (inferred from p2) + # Its size (n_x * n_y = 12) matches its parent p2's size, + # which differs from either individual dimension. + assert "p1" in xds.data_vars + npt.assert_array_almost_equal(xds["p1"].values, p1_all_arr) + + +def test_parameter_with_setpoints_has_control_size_mismatch_warns( + experiment: "Experiment", caplog: "pytest.LogCaptureFixture" +) -> None: + """Test that a warning is emitted when the inferred parameter has a + different data size than its parent parameter.""" + + mp_data = np.arange(10) + + mp = ManualParameter("mp", vals=Arrays(shape=(10,)), initial_value=mp_data) + p1 = ParameterWithSetpoints( + "p1", vals=Arrays(shape=(10,)), setpoints=(mp,), set_cmd=None + ) + p2 = _make_controlled_setpoints( + "p2", p1, vals=Arrays(shape=(10,)), setpoints=(mp,), set_cmd=None + ) + + p1(np.linspace(-1, 1, 10)) + p2(np.random.default_rng().standard_normal(10)) + + meas = Measurement() + meas.register_parameter(p2) + with meas.run() as ds: + ds.add_result((p2, p2())) + + # Build an xarray dataset and sub_dict with mismatched p1 data to + # exercise the warning path in _add_inferred_data_vars directly. + + raw_data = ds.dataset.get_parameter_data() + sub_dict = dict(raw_data["p2"]) + # Replace p1 with wrong-sized data (5 instead of 10) + sub_dict["p1"] = np.zeros(5) + + xr_dataset = xr.Dataset( + {"p2": (("mp",), sub_dict["p2"].ravel())}, + coords={"mp": sub_dict["mp"].ravel()}, + ) + + with caplog.at_level( + logging.WARNING, logger="qcodes.dataset.exporters.export_to_xarray" + ): + result = _add_inferred_data_vars(ds.dataset, "p2", sub_dict, xr_dataset) + + assert "p1" not in result.data_vars + assert any( + "Cannot add inferred parameter 'p1'" in msg and "'p2'" in msg + for msg in caplog.messages + )