Add 2d infeered test

Author: jenshnielsen

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