Merge pull request #868 from genematx/fix-dims

Keep Dimension Names

Author: danielballan

bluesky-tiled-plugins/bluesky_tiled_plugins/consolidators.py

@@ -140,6 +140,9 @@ def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor)
         self._num_rows: int = 0  # Number of rows in the Data Source (all rows, includung skips)
         self._seqnums_to_indices_map: dict[int, int] = {}
 
+        # Set the dimension names if provided
+        self.dims: tuple[str, ...] = tuple(data_desc.get("dims", ()))
+
     @classmethod
     def get_supported_mimetype(cls, sres):
         if sres["mimetype"] not in cls.supported_mimetypes:
@@ -241,6 +244,7 @@ def structure(self) -> ArrayStructure:
             data_type=self.data_type,
             shape=self.shape,
             chunks=self.chunks,
+            dims=self.dims if self.dims else None,
         )
 
     def consume_stream_datum(self, doc: StreamDatum):
@@ -300,22 +304,24 @@ def update_from_stream_resource(self, stream_resource: StreamResource):
 
         raise NotImplementedError("This method is not implemented in the base Consolidator class.")
 
-    def validate(self, adapters_by_mimetype=None, fix_errors=False):
+    def validate(self, adapters_by_mimetype=None, fix_errors=False) -> list[str]:
         """Validate the Consolidator's state against the expected structure"""
 
         # User-provided adapters take precedence over defaults.
         all_adapters_by_mimetype = collections.ChainMap((adapters_by_mimetype or {}), DEFAULT_ADAPTERS_BY_MIMETYPE)
         adapter_class = all_adapters_by_mimetype[self.mimetype]
 
         # Initialize adapter from uris and determine the structure
-        uris = [asset.data_uri for asset in self.assets if asset.parameter == "data_uris"]
+        uris = [asset.data_uri for asset in self.assets]
         structure = adapter_class.from_uris(*uris, **self.adapter_parameters()).structure()
+        notes = []
 
         if self.shape != structure.shape:
             if not fix_errors:
                 raise ValueError(f"Shape mismatch: {self.shape} != {structure.shape}")
             else:
-                warnings.warn(f"Fixing shape mismatch: {self.shape} -> {structure.shape}", stacklevel=2)
+                msg = f"Fixed shape mismatch: {self.shape} -> {structure.shape}"
+                warnings.warn(msg, stacklevel=2)
                 if self.join_method == "stack":
                     self._num_rows = structure.shape[0]
                     self.datum_shape = structure.shape[1:]
@@ -324,26 +330,53 @@ def validate(self, adapters_by_mimetype=None, fix_errors=False):
                     multiplier = 1 if structure.shape[0] % structure.chunks[0][0] else structure.chunks[0][0]
                     self._num_rows = structure.shape[0] // multiplier
                     self.datum_shape = (multiplier,) + structure.shape[1:]
+                notes.append(msg)
 
         if self.chunks != structure.chunks:
             if not fix_errors:
                 raise ValueError(f"Chunk shape mismatch: {self.chunks} != {structure.chunks}")
             else:
                 _chunk_shape = tuple(c[0] for c in structure.chunks)
-                warnings.warn(f"Fixing chunk shape mismatch: {self.chunk_shape} -> {_chunk_shape}", stacklevel=2)
+                msg = f"Fixed chunk shape mismatch: {self.chunk_shape} -> {_chunk_shape}"
+                warnings.warn(msg, stacklevel=2)
                 self.chunk_shape = _chunk_shape
+                notes.append(msg)
 
         if self.data_type != structure.data_type:
             if not fix_errors:
                 raise ValueError(f"dtype mismatch: {self.data_type} != {structure.data_type}")
             else:
-                warnings.warn(
-                    f"Fixing dtype mismatch: {self.data_type.to_numpy_dtype()} -> {structure.data_type.to_numpy_dtype()}",  # noqa
-                    stacklevel=2,
+                msg = (
+                    f"Fixed dtype mismatch: {self.data_type.to_numpy_dtype()} "
+                    f"-> {structure.data_type.to_numpy_dtype()}"
                 )
+                warnings.warn(msg, stacklevel=2)
                 self.data_type = structure.data_type
+                notes.append(msg)
+
+        if self.dims and (len(self.dims) != len(structure.shape)):
+            if not fix_errors:
+                raise ValueError(
+                    f"Number of dimension names mismatch for a "
+                    f"{len(structure.shape)}-dimensional array: {self.dims}"
+                )
+            else:
+                old_dims = self.dims
+                if len(old_dims) < len(structure.shape):
+                    self.dims = (
+                        ("time",)
+                        + old_dims
+                        + tuple(f"dim{i}" for i in range(len(old_dims) + 1, len(structure.shape)))
+                    )
+                else:
+                    self.dims = old_dims[: len(structure.shape)]
+                msg = f"Fixed dimension names: {old_dims} -> {self.dims}"
+                warnings.warn(msg, stacklevel=2)
+                notes.append(msg)
+
+        assert self.get_adapter() is not None, "Adapter can not be initialized"
 
-        assert self.get_adapter() is not None, "Adapter can not not initialized"
+        return notes
 
 
 class CSVConsolidator(ConsolidatorBase):

bluesky-tiled-plugins/bluesky_tiled_plugins/tiled_writer.py

@@ -198,6 +198,7 @@ def __init__(
         self._emitted: set[str] = set()  # UIDs of the StreamResource documents that have been emitted
         self._int_keys: set[str] = set()  # Names of internal data_keys
         self._ext_keys: set[str] = set()
+        self.notes: list[str] = []  # Human-readable notes about any modifications made to the documents
 
     def _convert_resource_to_stream_resource(self, doc: Union[Resource, StreamResource]) -> StreamResource:
         """Make changes to and return a shallow copy of StreamRsource dictionary adhering to the new structure.
@@ -338,6 +339,8 @@ def stop(self, doc: RunStop):
                     f"Cannot emit StreamDatum for {data_key} because the corresponding Datum document is missing."
                 )
 
+        doc["_run_normalizer_notes"] = self.notes or []  # Add notes about modifications to the stop document
+
         self.emit(DocumentNames.stop, doc)
 
     def descriptor(self, doc: EventDescriptor):
@@ -351,14 +354,14 @@ def descriptor(self, doc: EventDescriptor):
                 if f"_{name}" in doc["data_keys"].keys():
                     raise ValueError(f"Cannot rename {name} to _{name} because it already exists")
                 doc["data_keys"][f"_{name}"] = doc["data_keys"].pop(name)
-                for obj_data_keys_list in doc["object_keys"].values():
+                for obj_data_keys_list in doc.get("object_keys", {}).values():
                     if name in obj_data_keys_list:
                         obj_data_keys_list.remove(name)
                         obj_data_keys_list.append(f"_{name}")
 
         # Rename some fields (in-place) to match the current schema for the descriptor
         # Loop over all dictionaries that specify data_keys (both event data_keys or configuration data_keys)
-        conf_data_keys = (obj["data_keys"].values() for obj in doc["configuration"].values())
+        conf_data_keys = (obj["data_keys"].values() for obj in doc.get("configuration", {}).values())
         for data_keys_spec in itertools.chain(doc["data_keys"].values(), *conf_data_keys):
             # Determine numpy data type. From highest precedent to lowest:
             # 1. Try 'dtype_descr', optional, if present -- this is a structural dtype
@@ -376,8 +379,9 @@ def descriptor(self, doc: EventDescriptor):
             ):
                 data_keys_spec["dtype_numpy"] = dtype_numpy
 
-        # Ensure that all event data_keys have object_name assigned (for consistency)
-        for obj_name, data_keys_list in doc["object_keys"].items():
+        # Ensure that all event data_keys have object_name assigned, if known (for consistency)
+        # If "object_keys" are not present, do not reconstruct them -- they are optional
+        for obj_name, data_keys_list in doc.get("object_keys", {}).items():
             for key in data_keys_list:
                 doc["data_keys"][key]["object_name"] = obj_name
 
@@ -466,6 +470,7 @@ def datum(self, doc: Datum):
         if patch := self.patches.get("datum"):
             doc = patch(doc)
 
+        # Keep the Datum document in memory until it is referenced by an Event document
         self._datum_cache[doc["datum_id"]] = doc
 
     def datum_page(self, doc: DatumPage):
@@ -613,9 +618,6 @@ def stop(self, doc: RunStop):
         self.root_node.update_metadata(metadata={"stop": doc, **dict(self.root_node.metadata)}, drop_revision=True)
 
     def descriptor(self, doc: EventDescriptor):
-        if self.root_node is None:
-            raise RuntimeError("RunWriter is not properly initialized: no Start document has been recorded.")
-
         desc_name = doc["name"]  # Name of the descriptor/stream
         self.data_keys.update(doc.get("data_keys", {}))
 
@@ -684,6 +686,7 @@ def get_sres_node(self, sres_uid: str, desc_uid: Optional[str] = None) -> tuple[
             if not desc_uid:
                 raise RuntimeError("Descriptor uid must be specified to initialise a Stream Resource node")
 
+            # Define `full_data_key` as desc_name + _ + data_key to ensure uniqueness across streams
             sres_doc = self._stream_resource_cache[sres_uid]
             desc_node = self._desc_nodes[desc_uid]
             full_data_key = f"{desc_node.item['id']}_{sres_doc['data_key']}"  # desc_name + data_key

bluesky-tiled-plugins/tests/examples/external_assets.json

@@ -43,6 +43,11 @@
                         13,
                         17
                     ],
+                    "dims": [
+                        "time",
+                        "dim_x",
+                        "dim_y"
+                    ],
                     "external": "STREAM:",
                     "object_name": "det-obj1"
                 }

bluesky-tiled-plugins/tests/test_tiled_writer.py

@@ -382,6 +382,20 @@ def test_with_correct_sample_runs(client, batch_size, external_assets_folder, fn
         assert stream.read() is not None
 
 
+def test_dims_names(client, external_assets_folder):
+    tw = TiledWriter(client)
+
+    for item in render_templated_documents("external_assets.json", external_assets_folder):
+        if item["name"] == "start":
+            uid = item["doc"]["uid"]
+        tw(**item)
+
+    run = client[uid]
+
+    assert run["primary"]["det-key1"].structure().dims is None
+    assert run["primary"]["det-key2"].structure().dims == ("time", "dim_x", "dim_y")
+
+
 @pytest.mark.parametrize(
     "batch_size, expected_patch_shapes, expected_patch_offsets",
     [(1, (1, 1, 1), (0, 1, 2)), (2, (2, 1), (0, 2)), (5, (3,), (0,))],