Added extra snapshot data to data_writer

core_tools/data/data_writer.py

@@ -41,7 +41,7 @@ class _Action:
 
 
 class DataWriter:
-    def __init__(self, name, *args):
+    def __init__(self, name, *args, snapshot_data: dict[str, any] | None = None):
         self._measurement = Measurement(name, silent=True)
         self._actions = []
         self._set_params = []
@@ -55,7 +55,10 @@ class DataWriter:
                 self._add_data(arg)
             else:
                 raise TypeError(f"Unknown argument of type {type(arg)}")
-        self._measurement.add_snapshot('data_writer', {'message': 'Data written by data writer'})
+        if snapshot_data:
+            self._measurement.add_snapshot('data_writer', snapshot_data)
+        else:
+            self._measurement.add_snapshot('data_writer', {'message': 'Data written by data writer'})
 
     def _add_axis(self, axis):
         param = ManualParameter(axis.name, label=axis.label, unit=axis.unit)
@@ -100,13 +103,14 @@ class DataWriter:
             self._loop(iaction + 1, isetpoint)
 
 
-def write_data(name: str, *args):
+def write_data(name: str, *args, snapshot_data: dict[str, any] | None = None):
     '''
     Creates a dataset `name` using the specified Axis and Data.
 
     Args:
         name: name of the dataset.
         args: list of Axis and Data objects.
+        snapshot_data: Data to be added to the snapshot.
 
     Example:
         write_data(
@@ -131,7 +135,7 @@ def write_data(name: str, *args):
             Data('z', 'z', 'a.u., <array with shape(len(values_a), len(values_b), len(values_c)>),
         )
     '''
-    return DataWriter(name, *args).run()
+    return DataWriter(name, *args, snapshot_data=snapshot_data).run()
 
 
 if __name__ == "__main__":

examples/demo_station/data_writer_demo.py

@@ -11,6 +11,11 @@ write_data(
     'Demo_WriteData',
     Axis('x', 'x-array', 'a.u.', [1, 2, 3, 4]),
     Data('y', 'y-array', 'a.u.', [3, 1, 4, 6]),
+    snapshot_data={
+        "model": "some arbitrary data",
+        "param_a": 100,
+        "param_b": 10,
+        }
 )
 
 write_data(
@@ -19,7 +24,9 @@ write_data(
     Axis('e12', 'detuning', 'mv', np.linspace(-10, 10, 5)),
     Data('SD1', 'Sensor 1', 'mV', np.linspace(10, 20, 55).reshape((11, 5))),
     Data('SD2', 'Sensor 2', 'mV', np.linspace(0, -20, 55).reshape((11, 5))),
+    snapshot_data={
+        "model": "linear numpy data",
+        "param_a": 100,
+        "param_b": 10,
+        }
 )
-
-
-#%%
\ No newline at end of file