Merge pull request #704 from European-XFEL/feat/keydata-train-index

Add  KeyData.train_index_bounds

Author: takluyver

extra_data/keydata.py

@@ -536,6 +536,29 @@ def train_id_coordinates(self):
         ]
         return np.concatenate(chunks_trainids)
 
+    def train_index_bounds(self, labelled=False):
+        """Generate first and last indices of trains to use alongside data
+        from ``.ndarray()``.
+
+        If *labelled* is True, returns a pandas dataframe with columns start
+        and stop. Otherwise, returns a tuple of two NumPy arrays.
+        """
+
+        counts = self.data_counts(labelled)
+        start = counts.copy()
+
+        if labelled:
+            start.iloc[0] = 0
+            start.iloc[1:] = counts.cumsum().iloc[:-1]
+            stop = start + counts
+            import pandas as pd
+            return pd.concat(
+                [start.rename('start'), stop.rename('stop')], axis=1)
+        else:
+            start[0] = 0
+            start[1:] = counts.cumsum()[:-1]
+            return start, start + counts
+
     def xarray(self, extra_dims=None, roi=(), name=None, extra_coords=None):
         """Load this data as a labelled xarray array or dataset.
 

extra_data/tests/mockdata/xgm.py

@@ -31,7 +31,7 @@ class XGM(DeviceBase):
         ('pulseEnergy/crossUsed', 'f4', ()),
         ('pulseEnergy/gammaUsed', 'f4', ()),
         ('pulseEnergy/gmdError', 'i4', ()),
-        ('pulseEnergy/nummberOfBrunches', 'f4', ()),
+        ('pulseEnergy/nummberOfBunches', 'f4', ()),
         ('pulseEnergy/photonFlux', 'f4', ()),
         ('pulseEnergy/pressure', 'f4', ()),
         ('pulseEnergy/temperature', 'f4', ()),

extra_data/tests/test_keydata.py

@@ -248,6 +248,33 @@ def test_data_counts_missing_train(fxe_run_module_offset):
     np.testing.assert_array_equal(arr, 128)
 
 
+@pytest.mark.parametrize('labelled', [True, False])
+def test_train_index_bounds(mock_spb_raw_run, labelled):
+    run = RunDirectory(mock_spb_raw_run)
+
+    agipd_m0 = run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf', 'image.pulseId']
+    bounds = agipd_m0.train_index_bounds(labelled)
+
+    if labelled:
+        start, stop = bounds['start'], bounds['stop']
+    else:
+        start, stop = bounds
+
+    np.testing.assert_array_equal(start, np.arange(0, 4032+1, 64))
+    np.testing.assert_array_equal(stop, start+64)
+
+    xgm = run['SPB_XTD9_XGM/DOOCS/MAIN', 'pulseEnergy.photonFlux']
+    bounds = xgm.train_index_bounds(labelled)
+
+    if labelled:
+        start, stop = bounds['start'], bounds['stop']
+    else:
+        start, stop = bounds
+
+    np.testing.assert_array_equal(start, np.arange(len(start)))
+    np.testing.assert_array_equal(stop, start+1)
+
+
 def test_select_by(mock_spb_raw_run):
     run = RunDirectory(mock_spb_raw_run)
     am0 = run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf', 'image.data']