Expose units symbol & name on KeyData objects

extra_data/keydata.py

@@ -1,5 +1,6 @@
 from typing import List, Optional, Tuple
 
+import h5py
 import numpy as np
 
 from .exceptions import TrainIDError, NoDataError
@@ -106,6 +107,28 @@ def size_gb(self):
         """The size of the data in memory in gigabytes."""
         return self.nbytes / 1e9
 
+    @property
+    def units(self):
+        attrs = self.attributes()
+        base_unit = attrs.get('unitSymbol', None)
+        if base_unit is None:
+            return None
+
+        prefix = attrs.get('metricPrefixSymbol', '')
+        if prefix == 'u':
+            prefix = 'μ'  # We are not afraid of unicode
+        return prefix + base_unit
+
+    @property
+    def units_name(self):
+        attrs = self.attributes()
+        base_unit = attrs.get('unitName', None)
+        if base_unit is None:
+            return None
+
+        prefix = attrs.get('metricPrefixName', '')
+        return prefix + base_unit
+
     @property
     def source_file_paths(self):
         paths = []
@@ -129,6 +152,18 @@ def source_file_paths(self):
         from pathlib import Path
         return [Path(p) for p in paths]
 
+    def attributes(self):
+        dset = self.files[0].file[self.hdf5_data_path]
+        if len(dset.attrs) == 0 and dset.is_virtual:
+            # Virtual datasets were initially created without these attributes.
+            # Find a source file. Not using source_file_paths as it can give [].
+            _, filename, _, _ = dset.virtual_sources()[0]
+            # Not using FileAccess: no need for train or source lists.
+            f = h5py.File(filename, 'r')
+            dset = f[self.hdf5_data_path]
+
+        return dict(dset.attrs)
+
     def select_trains(self, trains):
         """Select a subset of trains in this data as a new :class:`KeyData` object.
 
@@ -144,8 +179,11 @@ def __getitem__(self, item):
 
     def _only_tids(self, tids):
         tids_arr = np.array(tids)
-        files = [f for f in self.files
-                 if f.has_train_ids(tids_arr, self.inc_suspect_trains)]
+        # Keep 1 file, even if 0 trains selected.
+        files = [
+            f for f in self.files
+            if f.has_train_ids(tids_arr, self.inc_suspect_trains)
+        ] or [self.files[0]]
 
         return KeyData(
             self.source,

extra_data/tests/mockdata/xgm.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from .base import DeviceBase
 
 class XGM(DeviceBase):
@@ -53,3 +55,16 @@ class XGM(DeviceBase):
         ('xTD', 'f4', (1000,)),
         ('yTD', 'f4', (1000,)),
     ]
+
+    def write_instrument(self, f):
+        super().write_instrument(f)
+
+        # Annotate intensityTD with some units to test retrieving them
+        # Karabo stores ASCII strings, assigning bytes is a shortcut to mimic that
+        ds = f[f'INSTRUMENT/{self.device_id}:output/data/intensityTD']
+        ds.attrs['metricPrefixEnum']= np.array([14], dtype=np.int32)
+        ds.attrs['metricPrefixName'] = b'micro'
+        ds.attrs['metricPrefixSymbol'] = b'u'
+        ds.attrs['unitEnum'] = np.array([15], dtype=np.int32)
+        ds.attrs['unitName'] = b'joule'
+        ds.attrs['unitSymbol'] = b'J'

extra_data/tests/test_keydata.py

@@ -49,7 +49,7 @@ def test_select_trains(mock_spb_raw_run):
     # Empty selection
     sel2 = xgm_beam_x[80:]
     assert sel2.shape == (0,)
-    assert len(sel2.files) == 0
+    assert len(sel2.files) == 1
     assert sel2.xarray().shape == (0,)
 
     # Single train
@@ -339,3 +339,24 @@ def test_file_no_trains(run_with_file_no_trains):
     run = RunDirectory(run_with_file_no_trains)
     xpos = run['SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos'].ndarray()
     assert xpos.shape == (64,)
+
+
+def test_attributes(mock_sa3_control_data):
+    run = H5File(mock_sa3_control_data)
+    xgm_intensity = run['SA3_XTD10_XGM/XGM/DOOCS:output', 'data.intensityTD']
+    attrs = xgm_intensity.attributes()
+
+    assert isinstance(attrs, dict)
+    assert attrs['metricPrefixName'] == 'micro'
+    assert attrs['unitSymbol'] == 'J'
+
+
+def test_units(mock_sa3_control_data):
+    run = H5File(mock_sa3_control_data)
+    xgm_intensity = run['SA3_XTD10_XGM/XGM/DOOCS:output', 'data.intensityTD']
+
+    assert xgm_intensity.units == 'μJ'
+    assert xgm_intensity.units_name == 'microjoule'
+
+    # Check that it still works after selecting 0 trains
+    assert xgm_intensity.select_trains(np.s_[:0]).units == 'μJ'

extra_data/tests/test_voview.py

@@ -63,6 +63,8 @@ def test_use_voview(mock_spb_raw_run, tmp_path):
     assert {p.name for p in xgm_intens[:30].source_file_paths} == {
         'RAW-R0238-DA01-S00000.h5'
     }
+    assert xgm_intens.units == 'μJ'
+    assert xgm_intens.units_name == 'microjoule'