Expand support for RUN data

extra_data/keydata.py

@@ -105,7 +105,7 @@ class KeyData:
     """
     def __init__(
             self, source, key, *, train_ids, files, section, dtype, eshape,
-            inc_suspect_trains=True,
+            is_single_run, inc_suspect_trains=True,
     ):
         self.source = source
         self.key = key
@@ -115,6 +115,7 @@ def __init__(
         self.dtype = dtype
         self.entry_shape = eshape
         self.ndim = len(eshape) + 1
+        self.is_single_run = is_single_run
         self.inc_suspect_trains = inc_suspect_trains
 
     def _find_chunks(self):
@@ -343,6 +344,7 @@ def _only_tids(self, tids, files=None):
             section=self.section,
             dtype=self.dtype,
             eshape=self.entry_shape,
+            is_single_run=self.is_single_run,
             inc_suspect_trains=self.inc_suspect_trains,
         )
 
@@ -472,6 +474,29 @@ def as_single_value(self, rtol=1e-5, atol=0.0, reduce_by=None):
 
         return value
 
+    def run_value(self, allow_multi_run=False):
+        """Get the RUN value for this key if it exists.
+
+        This method is intended for use with data from a single run. If you
+        combine data from multiple runs, it will raise MultiRunError.
+
+        Returns the RUN parameter value corresponding to this key.
+        """
+
+        from .sourcedata import SourceData  # Prevent cyclic import.
+
+        # Construct minimal SourceData object to obtain RUN value.
+        return SourceData(
+            self.source,
+            sel_keys=None,
+            train_ids=self.train_ids,
+            files=self.files,
+            section=self.section,
+            canonical_name=self.source,
+            is_single_run=self.is_single_run,
+            inc_suspect_trains=self.inc_suspect_trains
+        ).run_value(self.key, allow_multi_run=allow_multi_run)
+
     # Getting data as different kinds of array: -------------------------------
 
     def ndarray(self, roi=(), out=None):

extra_data/sourcedata.py

@@ -118,6 +118,7 @@ def __getitem__(self, key):
             section=self.section,
             dtype=ds0.dtype,
             eshape=ds0.shape[1:],
+            is_single_run=self.is_single_run,
             inc_suspect_trains=self.inc_suspect_trains,
         )
 
@@ -473,12 +474,12 @@ def run_value(self, key, *, allow_multi_run=False):
 
         Returns the RUN parameter value corresponding to the *key* argument.
         """
+
         if not (self.is_single_run or allow_multi_run):
             raise MultiRunError()
 
-        if not self.is_control:
-            raise ValueError('Only CONTROL sources have run values, '
-                             f'{self.source} is a/an {self.section} source')
+        if self.source not in self.files[0].file['RUN']:
+            raise ValueError(f'{self.source} has no RUN values')
 
         # Arbitrary file - should be the same across a run
         ds = self.files[0].file['RUN'][self.source].get(key.replace('.', '/'))
@@ -493,6 +494,29 @@ def run_value(self, key, *, allow_multi_run=False):
             return val.decode('utf-8', 'surrogateescape')
         return val
 
+    def run_keys(self, inc_timestamps=True):
+        """Get a set of RUN keys for this source.
+
+        Unlike :meth:`keys`, RUN keys are not affected by selection.
+        """
+
+        if not self.is_single_run:
+            raise MultiRunError()
+
+        if self.source not in self.files[0].file['RUN']:
+            raise ValueError(f'{self.source} has no RUN values')
+
+        res = set()
+        def visitor(path, obj):
+            if isinstance(obj, h5py.Dataset):
+                res.add(path.replace('/', '.'))
+
+        # Arbitrary file - should be the same across a run
+        self.files[0].file['RUN'][self.source].visititems(visitor)
+        if not inc_timestamps:
+            return {k[:-6] for k in res if k.endswith('.value')}
+        return res
+
     def run_values(self, inc_timestamps=True):
         """Get a dict of all RUN values for this source
 
@@ -501,9 +525,8 @@ def run_values(self, inc_timestamps=True):
         if not self.is_single_run:
             raise MultiRunError()
 
-        if not self.is_control:
-            raise ValueError('Only CONTROL sources have run values, '
-                             f'{self.source} is a/an {self.section} source')
+        if self.source not in self.files[0].file['RUN']:
+            raise ValueError(f'{self.source} has no RUN values')
 
         res = {}
         def visitor(path, obj):

extra_data/tests/test_keydata.py

@@ -363,6 +363,17 @@ def test_single_value(mock_sa3_control_data, monkeypatch):
     np.testing.assert_equal(intensity.as_single_value(rtol=1), np.median(data))
 
 
+def test_run_value(mock_sa3_control_data):
+    f = H5File(mock_sa3_control_data)
+
+    flux = f['SA3_XTD10_XGM/XGM/DOOCS', 'pulseEnergy.photonFlux']
+    assert flux.run_value() == 0.0
+
+    imager = f['SA3_XTD10_IMGFEL/CAM/BEAMVIEW:daqOutput', 'data.image.pixels']
+    with pytest.raises(ValueError):
+        assert imager.run_value()
+
+
 def test_ndarray_out(mock_spb_raw_run):
     f = RunDirectory(mock_spb_raw_run)
     cam = f['SPB_IRU_CAM/CAM/SIDEMIC:daqOutput', 'data.image.dims']

extra_data/tests/test_sourcedata.py

@@ -186,8 +186,11 @@ def test_run_value(mock_spb_raw_run):
     assert 'pulseEnergy.conversion' in values_dict
     assert 'pulseEnergy.conversion.timestamp' not in values_dict
 
+    assert xgm.run_keys() == xgm.keys() | {'classId.timestamp', 'classId.value'}
+    assert xgm.run_keys(False) == xgm.keys(False) | {'classId'}
+
     with pytest.raises(ValueError):
-        # no run values for instrument sources
+        # no run values for AGIPD instrument data.
         am0.run_values()