Merge pull request #6487 from samantha-ho/samanthaho/registered-param…

…eter-snapshot

Samanthaho/registered parameter snapshot

docs/changes/newsfragments/6487.improved

@@ -0,0 +1 @@
+Parameters registered in a qcodes Measurement are now snapshotted and stored in the resulting dataset under `dataset.snapshot["parameters"]`

src/qcodes/dataset/data_set.py

@@ -329,7 +329,7 @@ def prepare(
         parent_datasets: Sequence[Mapping[Any, Any]] = (),
         write_in_background: bool = False,
     ) -> None:
-        self.add_snapshot(json.dumps({"station": snapshot}, cls=NumpyJSONEncoder))
+        self.add_snapshot(json.dumps(snapshot, cls=NumpyJSONEncoder))
 
         if interdeps == InterDependencies_():
             raise RuntimeError("No parameters supplied")

src/qcodes/dataset/data_set_in_memory.py

@@ -421,7 +421,7 @@ def prepare(
         if not self.pristine:
             raise RuntimeError("Cannot prepare a dataset that is not pristine.")
 
-        self.add_snapshot(json.dumps({"station": snapshot}, cls=NumpyJSONEncoder))
+        self.add_snapshot(json.dumps(snapshot, cls=NumpyJSONEncoder))
 
         if interdeps == InterDependencies_():
             raise RuntimeError("No parameters supplied")

src/qcodes/dataset/measurements.py

@@ -549,6 +549,7 @@ def __init__(
         in_memory_cache: bool | None = None,
         dataset_class: DataSetType = DataSetType.DataSet,
         parent_span: trace.Span | None = None,
+        registered_parameters: Sequence[ParameterBase] | None = None,
     ) -> None:
         if in_memory_cache is None:
             in_memory_cache = qc.config.dataset.in_memory_cache
@@ -577,6 +578,7 @@ def __init__(
         self._in_memory_cache = in_memory_cache
         self._parent_span = parent_span
         self.ds: DataSetProtocol
+        self._registered_parameters = registered_parameters
 
     @staticmethod
     def _calculate_write_period(
@@ -661,9 +663,15 @@ def __enter__(self) -> DataSaver:
             station = self.station
 
         if station is not None:
-            snapshot = station.snapshot()
+            snapshot = {"station": station.snapshot()}
         else:
             snapshot = {}
+        if self._registered_parameters is not None:
+            parameter_snapshot = {
+                param.short_name: param.snapshot()
+                for param in self._registered_parameters
+            }
+            snapshot["parameters"] = parameter_snapshot
 
         self.ds.prepare(
             snapshot=snapshot,
@@ -803,6 +811,7 @@ def __init__(
         self._shapes: Shapes | None = None
         self._parent_datasets: list[dict[str, str]] = []
         self._extra_log_info: str = ""
+        self._registered_parameters: list[ParameterBase] = []
 
     @property
     def parameters(self) -> dict[str, ParamSpecBase]:
@@ -926,7 +935,6 @@ def register_parameter(
                 f"Can not register object of type {type(parameter)}. Can only "
                 "register a QCoDeS Parameter."
             )
-
         paramtype = self._infer_paramtype(parameter, paramtype)
         # default to numeric
         if paramtype is None:
@@ -981,6 +989,7 @@ def register_parameter(
             raise RuntimeError(
                 f"Does not know how to register a parameter of type {type(parameter)}"
             )
+        self._registered_parameters.append(parameter)
 
         return self
 
@@ -1034,6 +1043,7 @@ def _register_parameter(
         setpoints: setpoints_type | None,
         basis: setpoints_type | None,
         paramtype: str,
+        metadata: dict[str, Any] | None = None,
     ) -> T:
         """
         Update the interdependencies object with a new group
@@ -1278,23 +1288,37 @@ def unregister_parameter(self, parameter: setpoints_type) -> None:
         running this measurement
         """
         if isinstance(parameter, ParameterBase):
-            param = str_or_register_name(parameter)
+            param_name = str_or_register_name(parameter)
         elif isinstance(parameter, str):
-            param = parameter
+            param_name = parameter
         else:
             raise ValueError(
                 "Wrong input type. Must be a QCoDeS parameter or"
                 " the name (a string) of a parameter."
             )
 
         try:
-            paramspec: ParamSpecBase = self._interdeps[param]
+            paramspec: ParamSpecBase = self._interdeps[param_name]
         except KeyError:
             return
 
         self._interdeps = self._interdeps.remove(paramspec)
 
-        log.info(f"Removed {param} from Measurement.")
+        # Must follow interdeps removal, because interdeps removal may error
+        if isinstance(parameter, ParameterBase):
+            try:
+                self._registered_parameters.remove(parameter)
+            except ValueError:
+                return
+        elif isinstance(parameter, str):
+            with_parameters_removed = [
+                param
+                for param in self._registered_parameters
+                if parameter not in (param.name, param.register_name)
+            ]
+            self._registered_parameters = with_parameters_removed
+
+        log.info(f"Removed {param_name} from Measurement.")
 
     def add_before_run(self: T, func: Callable[..., Any], args: Sequence[Any]) -> T:
         """
@@ -1412,6 +1436,7 @@ def run(
             in_memory_cache=in_memory_cache,
             dataset_class=dataset_class,
             parent_span=parent_span,
+            registered_parameters=self._registered_parameters,
         )
 
 

tests/dataset/test_snapshot.py

@@ -43,7 +43,9 @@ def test_station_snapshot_during_measurement(
 
     measurement.register_parameter(dac.ch1)
     measurement.register_parameter(dmm.v1, setpoints=[dac.ch1])
-
+    snapshot_of_parameters = {
+        parameter.short_name: parameter.snapshot() for parameter in (dac.ch1, dmm.v1)
+    }
     with measurement.run() as data_saver:
         data_saver.add_result((dac.ch1, 7), (dmm.v1, 5))
 
@@ -53,7 +55,10 @@ def test_station_snapshot_during_measurement(
     json_snapshot_from_dataset = data_saver.dataset.get_metadata("snapshot")  # type: ignore[attr-defined]
     snapshot_from_dataset = json.loads(json_snapshot_from_dataset)
 
-    expected_snapshot = {"station": snapshot_of_station}
+    expected_snapshot = {
+        "station": snapshot_of_station,
+        "parameters": snapshot_of_parameters,
+    }
     assert expected_snapshot == snapshot_from_dataset
 
     # 2. Test `snapshot_raw` property