Merge pull request #7502 from jenshnielsen/jenshnielsen/fix_cache_with_inferred_params_2

Write NaN values to cache when data is missing

Author: jenshnielsen

docs/changes/newsfragments/7502.improved

@@ -0,0 +1 @@
+Fixed a bug where omitting data for one or more variables could result in an inconsistent dataset cache. Missing data is now filled with appropriate empty values (0, "" or NaN depending on the data type)

src/qcodes/dataset/data_set_cache.py

@@ -325,7 +325,7 @@ def _merge_data(
     shape: tuple[int, ...] | None,
     single_tree_write_status: int | None,
     meas_parameter: str,
-) -> tuple[dict[str, npt.NDArray], int | None]:
+) -> tuple[dict[str, npt.NDArray], int]:
     subtree_merged_data = {}
     subtree_parameters = existing_data.keys()
 
@@ -335,20 +335,19 @@ def _merge_data(
             "The following keys were unexpected: "
             f"{set(new_data.keys() - existing_data.keys())}"
         )
-
-    new_write_status: int | None
-    single_param_merged_data, new_write_status = _merge_data_single_param(
+    single_param_merged_data, data_written = _merge_data_single_param(
         existing_data.get(meas_parameter),
         new_data.get(meas_parameter),
         shape,
         single_tree_write_status,
     )
+    new_write_status = data_written if data_written is not None else 0
     if single_param_merged_data is not None:
         subtree_merged_data[meas_parameter] = single_param_merged_data
 
     for subtree_param in subtree_parameters:
         if subtree_param != meas_parameter:
-            single_param_merged_data, new_write_status = _merge_data_single_param(
+            single_param_merged_data, data_written = _merge_data_single_param(
                 existing_data.get(subtree_param),
                 new_data.get(subtree_param),
                 shape,
@@ -357,6 +356,9 @@ def _merge_data(
             if single_param_merged_data is not None:
                 subtree_merged_data[subtree_param] = single_param_merged_data
 
+            if data_written is not None and data_written > new_write_status:
+                new_write_status = data_written
+
     return subtree_merged_data, new_write_status
 
 
@@ -373,22 +375,34 @@ def _merge_data_single_param(
         (merged_data, new_write_status) = _insert_into_data_dict(
             existing_values, new_values, single_tree_write_status, shape=shape
         )
-    elif new_values is not None:
+    elif new_values is not None or shape is not None:
         (merged_data, new_write_status) = _create_new_data_dict(new_values, shape)
     elif existing_values is not None:
         merged_data = existing_values
         new_write_status = single_tree_write_status
+    elif shape is None and new_values is None:
+        merged_data = existing_values
+        new_write_status = single_tree_write_status
     else:
         merged_data = None
         new_write_status = None
     return merged_data, new_write_status
 
 
 def _create_new_data_dict(
-    new_values: npt.NDArray, shape: tuple[int, ...] | None
-) -> tuple[npt.NDArray, int]:
-    if shape is None:
+    new_values: npt.NDArray | None, shape: tuple[int, ...] | None
+) -> tuple[npt.NDArray, int | None]:
+    if shape is None and new_values is None:
+        raise RuntimeError("Cannot create new data dict without new values")
+    elif shape is None:
+        assert new_values is not None
         return new_values, new_values.size
+    elif new_values is None:
+        # we don't know the datatype so use float which can hold NaN
+        # since that is the most common?
+        data = np.zeros(shape)
+        data[:] = np.nan
+        return data, None
     elif new_values.size > 0:
         n_values = new_values.size
         data = np.zeros(shape, dtype=new_values.dtype)

tests/dataset/test_dataset_in_memory.py

@@ -3,23 +3,29 @@
 import re
 import shutil
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import hypothesis.strategies as hst
 import numpy as np
 import pytest
 import xarray as xr
+from deepdiff import DeepDiff  # type: ignore[import-untyped]
 from hypothesis import HealthCheck, given, settings
 from numpy.testing import assert_almost_equal
 
 import qcodes
-from qcodes.dataset import load_by_id, load_by_run_spec
+from qcodes.dataset import Measurement, load_by_id, load_by_run_spec
 from qcodes.dataset.data_set_in_memory import DataSetInMem, load_from_file
 from qcodes.dataset.data_set_protocol import DataSetType
 from qcodes.dataset.descriptions.dependencies import InterDependencies_
 from qcodes.dataset.descriptions.param_spec import ParamSpecBase
 from qcodes.dataset.sqlite.connection import AtomicConnection, atomic_transaction
+from qcodes.parameters import ManualParameter, Parameter
 from qcodes.station import Station
 
+if TYPE_CHECKING:
+    from qcodes.dataset.experiment_container import Experiment
+
 
 def test_dataset_in_memory_reload_from_db(
     meas_with_registered_param, DMM, DAC, tmp_path
@@ -676,3 +682,50 @@ def test_load_from_db_dataset_moved(
             not in new_xr_ds.attrs
         )
         assert new_xr_ds.attrs["metadata_added_after_set_new_netcdf_location"] == 6969
+
+
+@pytest.mark.parametrize("include_inferred_data", [True, False])
+def test_dataset_in_mem_with_inferred_parameters(
+    experiment: "Experiment", include_inferred_data: bool
+) -> None:
+    inferred1 = ManualParameter("inferred1", initial_value=0.0)
+    inferred2 = ManualParameter("inferred2", initial_value=0.0)
+    control1 = ManualParameter("control1", initial_value=0.0)
+    control2 = ManualParameter("control2", initial_value=0.0)
+    dependent = Parameter("dependent", get_cmd=lambda: control1(), set_cmd=False)
+    meas = Measurement(exp=experiment, name="via Measurement")
+
+    meas.register_parameter(control1)
+    meas.register_parameter(control2)
+    meas.register_parameter(inferred1, basis=(control1, control2))
+    meas.register_parameter(inferred2, basis=(control1, control2))
+    meas.register_parameter(dependent, setpoints=(control1, control2))
+    meas.set_shapes({dependent.register_name: (11, 11)})
+    with meas.run() as datasaver:
+        for i in range(11):
+            for j in range(11):
+                control1(float(i))
+                control2(float(j))
+                if include_inferred_data:
+                    datasaver.add_result(
+                        (inferred1, inferred1()),
+                        (inferred2, inferred2()),
+                        (control1, control1()),
+                        (control2, control2()),
+                        (dependent, dependent()),
+                    )
+                else:
+                    datasaver.add_result(
+                        (control1, control1()),
+                        (control2, control2()),
+                        (dependent, dependent()),
+                    )
+        ds = datasaver.dataset
+
+    param_data = ds.get_parameter_data()
+    cache_data = ds.cache.data()
+
+    assert set(param_data.keys()) == set(cache_data.keys())
+    assert set(param_data["dependent"].keys()) == set(cache_data["dependent"].keys())
+
+    assert DeepDiff(param_data, cache_data, ignore_nan_inequality=True) == {}