From 7e917601eeac1b774246198ed37e80b193c0b0f7 Mon Sep 17 00:00:00 2001 From: yuk Date: Mon, 17 Oct 2022 11:31:13 +0800 Subject: [PATCH] Cleanup grammar, styles and code Signed-off-by: Egor Savkin Revise metadata logic Raise on absent keys, raise on wrong types, do not allow metadata on dataset itself. Make tests more real-world-like. Signed-off-by: Egor Savkin Add HDF5 metadata unit tests Signed-off-by: Egor Savkin Add HDF5 attributes support to release notes Signed-off-by: yuk hdf5: Warn when attaching metadata to absent key Signed-off-by: yuk hdf5: Add demo for saving attributes Signed-off-by: yuk hdf5: Allow saving custom attributes Signed-off-by: yuk --- RELEASE_NOTES.rst | 5 +- .../no_hardware/repository/hdf5_attributes.py | 15 +++++ artiq/language/environment.py | 15 +++++ artiq/master/worker_db.py | 15 ++++- artiq/test/test_hdf5_attributes.py | 59 +++++++++++++++++++ 5 files changed, 106 insertions(+), 3 deletions(-) create mode 100644 artiq/examples/no_hardware/repository/hdf5_attributes.py create mode 100644 artiq/test/test_hdf5_attributes.py diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index e3bab5ca1e..689851f233 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -7,6 +7,7 @@ ARTIQ-9 (Unreleased) -------------------- * Zadig Driver Installer was added to the MSYS2 offline installer. +* HDF5 attributes can be attached to datasets using ``set_dataset_metadata()``. ARTIQ-8 ------- @@ -49,7 +50,7 @@ Highlights: - Hotkeys now organize experiment windows in the order they were last interacted with: + CTRL+SHIFT+T tiles experiment windows + CTRL+SHIFT+C cascades experiment windows - - By enabling the ``quickstyle`` option, ``EnumerationValue`` entry widgets can now alternatively display + - By enabling the ``quickstyle`` option, ``EnumerationValue`` entry widgets can now alternatively display its choices as buttons that submit the experiment on click. * Datasets can now be associated with units and scale factors, and displayed accordingly in the dashboard including applets, like widgets such as ``NumberValue`` already did in earlier ARTIQ versions. @@ -92,7 +93,7 @@ Accesses to the data argument should be replaced as below: data[key][0] ==> persist[key] data[key][1] ==> value[key] -* The ``ndecimals`` parameter in ``NumberValue`` and ``Scannable`` has been renamed to ``precision``. +* The ``ndecimals`` parameter in ``NumberValue`` and ``Scannable`` has been renamed to ``precision``. Parameters after and including ``scale`` in both constructors are now keyword-only. Refer to the updated ``no_hardware/arguments_demo.py`` example for current usage. * Almazny v1.2 is incompatible with the legacy versions and is the default. diff --git a/artiq/examples/no_hardware/repository/hdf5_attributes.py b/artiq/examples/no_hardware/repository/hdf5_attributes.py new file mode 100644 index 0000000000..a631f5afff --- /dev/null +++ b/artiq/examples/no_hardware/repository/hdf5_attributes.py @@ -0,0 +1,15 @@ +import numpy as np + +from artiq.experiment import * + + +class HDF5Attributes(EnvExperiment): + """Archive data to HDF5 with attributes""" + + def run(self): + dummy = np.empty(20) + dummy.fill(np.nan) + self.set_dataset("dummy", dummy, + broadcast=True, archive=True) + self.set_dataset_metadata("dummy", "k1", "v1") + self.set_dataset_metadata("dummy", "k2", "v2") diff --git a/artiq/language/environment.py b/artiq/language/environment.py index a80701c782..e60e75a1c6 100644 --- a/artiq/language/environment.py +++ b/artiq/language/environment.py @@ -444,6 +444,21 @@ def append_to_dataset(self, key, value): efficiently as incremental modifications in broadcast mode.""" self.__dataset_mgr.append_to(key, value) + @rpc(flags={"async"}) + def set_dataset_metadata(self, key, metadata_key, metadata_value): + """Attach metadata to the dataset. + + The metadata is saved as HDF5 attributes if there was a call to + ``set_dataset(..., archive=True)`` with the same key. + + :param key: The already existing dataset, to which you want to attach the metadata. + If absent, KeyError will be raised. + :param metadata_key: The metadata key of type string. If already exists, rewrites the metadata. + :param metadata_value: Value to be attached to ``metadata_key``. Can be any valid HDF5 datatype. + See HDF5 documentation for additional information. + """ + self.__dataset_mgr.set_metadata(key, metadata_key, metadata_value) + def get_dataset(self, key, default=NoDefault, archive=True): """Returns the contents of a dataset. diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py index af316e4cfe..393303abd2 100644 --- a/artiq/master/worker_db.py +++ b/artiq/master/worker_db.py @@ -120,6 +120,7 @@ class DatasetManager: def __init__(self, ddb): self._broadcaster = Notifier(dict()) self.local = dict() + self.hdf5_attributes = dict() self.archive = dict() self.metadata = dict() @@ -142,7 +143,7 @@ def set(self, key, value, metadata, broadcast, persist, archive): self.local[key] = value elif key in self.local: del self.local[key] - + self.metadata[key] = metadata def _get_mutation_target(self, key): @@ -184,12 +185,24 @@ def get_metadata(self, key): return self.metadata[key] return self.ddb.get_metadata(key) + def set_metadata(self, key, metadata_key, metadata_value): + if key not in self.local: + raise KeyError(f"Dataset '{key}' does not exist.") + if key not in self.hdf5_attributes: + self.hdf5_attributes[key] = dict() + self.hdf5_attributes[key][metadata_key] = metadata_value + def write_hdf5(self, f): datasets_group = f.create_group("datasets") for k, v in self.local.items(): m = self.metadata.get(k, {}) _write(datasets_group, k, v, m) + for k, attrs in self.hdf5_attributes.items(): + assert k in datasets_group + for attr_k, attr_v in attrs.items(): + datasets_group[k].attrs[attr_k] = attr_v + archive_group = f.create_group("archive") for k, v in self.archive.items(): m = self.metadata.get(k, {}) diff --git a/artiq/test/test_hdf5_attributes.py b/artiq/test/test_hdf5_attributes.py new file mode 100644 index 0000000000..e7c0f691b3 --- /dev/null +++ b/artiq/test/test_hdf5_attributes.py @@ -0,0 +1,59 @@ +import unittest +import io +import numpy as np +import h5py + +from artiq.experiment import * +from artiq.test.hardware_testbench import ExperimentCase + + +class HDF5Attributes(EnvExperiment): + """Archive data to HDF5 with attributes""" + + def run(self): + # Attach attributes metadata to the HDF5 key + # The key should exist in the resulting HDF5 file (archive=True). + self.set_dataset("dummy", np.full(20, np.nan), broadcast=True, archive=True) + self.set_dataset_metadata("dummy", "k1", "v1") + self.set_dataset_metadata("dummy", "k2", "v2") + + +class TestHDF5Attributes(ExperimentCase): + def setUp(self): + super().setUp() + self.exp = self.execute(HDF5Attributes) + self.dump() + + def dump(self): + self.bio = io.BytesIO() + with h5py.File(self.bio, "w") as f: + self.dataset_mgr.write_hdf5(f) + + self.bio.seek(0) + self.h5file = h5py.File(self.bio, "r") + self.datasets = self.h5file.get("datasets") + + def test_dataset_metadata(self): + self.assertEqual(self.datasets["dummy"].attrs, {"k1": "v1", "k2": "v2"}) + self.assertTrue(np.all((self.datasets["dummy"], np.full(20, np.nan)))) + + def test_write_none(self): + with self.assertRaises(KeyError): + self.exp.set_dataset_metadata(None, "test", "none") + self.exp.set_dataset_metadata("dummy", None, "none") + with self.assertRaises(TypeError): + self.dump() + + def test_write_absent(self): + with self.assertRaises(KeyError): + self.exp.set_dataset_metadata("absent", "test", "absent") + + def test_rewrite(self): + self.exp.set_dataset_metadata("dummy", "k2", "rewrite") + self.dump() + self.assertEqual(self.datasets["dummy"].attrs, {"k1": "v1", "k2": "rewrite"}) + + def test_non_archive(self): + self.exp.set_dataset("non_archive", np.full(30, np.nan), broadcast=True, archive=False) + with self.assertRaises(KeyError): + self.exp.set_dataset_metadata("non_archive", "k1", "v1")