Repository: European-XFEL/EXtra-data

Issue:
train_ids in KeyData
the `train_ids` attribute of the `KeyData` object is a simple copy of the train ids contained in the `DataCollection`. That leads to some inconsistencies. For example calling `.shape` will return a shape where the trainId axis match the actual count of data available in the files for the source, but calling `.train_ids` will contains train ids with and without data for this key.

One way to go around this if for example using xarray: `data.xarray().trainIds` will return only the train ids containing data, but that'd be ideal if it wasn't necessary to load the data to find which trainIds have or do not have data.

Would it be sensible to have the `train_ids` filtered (either at init or when the property is called) and only contains the train ids that actually have data for the key?

Discussion:

Patch:
diff --git a/extra_data/keydata.py b/extra_data/keydata.py
index 858fb96..5e5573e 100644
--- a/extra_data/keydata.py
+++ b/extra_data/keydata.py
@@ -95,6 +95,27 @@ class KeyData:
     def __getitem__(self, item):
         return self.select_trains(item)
 
+    def data_counts(self):
+        """Get a count of data points in each train.
+
+        Returns a pandas series with an index of train IDs.
+        """
+        import pandas as pd
+
+        seq_series = []
+
+        for f in self.files:
+            if self.section == 'CONTROL':
+                counts = np.ones(len(f.train_ids), dtype=np.uint64)
+            else:
+                _, counts = f.get_index(self.source, self._key_group)
+            seq_series.append(pd.Series(counts, index=f.train_ids))
+
+        ser = pd.concat(sorted(seq_series, key=lambda s: s.index[0]))
+        # Select out only the train IDs of interest
+        train_ids = ser.index.intersection(self.train_ids)
+        return ser.loc[train_ids]
+
     # Getting data as different kinds of array: -------------------------------
 
     def ndarray(self, roi=()):
diff --git a/extra_data/reader.py b/extra_data/reader.py
index fe1f038..c22a210 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -409,23 +409,7 @@ class DataCollection:
         key: str
             Key of parameter within that device, e.g. "image.data".
         """
-        import pandas as pd
-
-        self._check_field(source, key)
-        seq_series = []
-
-        for f in self._source_index[source]:
-            if source in self.control_sources:
-                counts = np.ones(len(f.train_ids), dtype=np.uint64)
-            else:
-                group = key.partition('.')[0]
-                _, counts = f.get_index(source, group)
-            seq_series.append(pd.Series(counts, index=f.train_ids))
-
-        ser = pd.concat(sorted(seq_series, key=lambda s: s.index[0]))
-        # Select out only the train IDs of interest
-        train_ids = ser.index.intersection(self.train_ids)
-        return ser.loc[train_ids]
+        return self._get_key_data(source, key).data_counts()
 
     def get_series(self, source, key):
         """Return a pandas Series for a particular data field.


Test patch:
diff --git a/extra_data/tests/make_examples.py b/extra_data/tests/make_examples.py
index b3fde71..3d7625a 100644
--- a/extra_data/tests/make_examples.py
+++ b/extra_data/tests/make_examples.py
@@ -269,8 +269,7 @@ def make_reduced_spb_run(dir_path, raw=True, rng=None, format_version='0.5'):
             AGIPDModule('SPB_DET_AGIPD1M-1/DET/{}CH0'.format(modno), raw=raw,
                          frames_per_train=frame_counts)
             ], ntrains=64, chunksize=32, format_version=format_version)
-    if not raw:
-        return
+
     write_file(osp.join(dir_path, '{}-R0238-DA01-S00000.h5'.format(prefix)),
                [ XGM('SA1_XTD2_XGM/DOOCS/MAIN'),
                  XGM('SPB_XTD9_XGM/DOOCS/MAIN'),
diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py
index dc2657b..7edf40f 100644
--- a/extra_data/tests/test_keydata.py
+++ b/extra_data/tests/test_keydata.py
@@ -84,3 +84,23 @@ def test_get_train(mock_spb_raw_run):
 
     with pytest.raises(IndexError):
         xgm_beam_x.train_from_index(9999)
+
+
+def test_data_counts(mock_reduced_spb_proc_run):
+    run = RunDirectory(mock_reduced_spb_proc_run)
+
+    # control data
+    xgm_beam_x = run['SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos.value']
+    count = xgm_beam_x.data_counts()
+    assert count.index.tolist() == xgm_beam_x.train_ids
+    assert (count.values == 1).all()
+
+    # intrument data
+    camera = run['SPB_IRU_CAM/CAM/SIDEMIC:daqOutput', 'data.image.pixels']
+    count = camera.data_counts()
+    assert count.index.tolist() == camera.train_ids
+
+    mod = run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf', 'image.data']
+    count = mod.data_counts()
+    assert count.index.tolist() == mod.train_ids
+    assert count.values.sum() == mod.shape[0]
diff --git a/extra_data/tests/test_reader_mockdata.py b/extra_data/tests/test_reader_mockdata.py
index 730e6d0..1000570 100644
--- a/extra_data/tests/test_reader_mockdata.py
+++ b/extra_data/tests/test_reader_mockdata.py
@@ -636,3 +636,10 @@ def test_permission():
 def test_empty_file_info(mock_empty_file, capsys):
     f = H5File(mock_empty_file)
     f.info()  # smoke test
+
+
+def test_get_data_counts(mock_spb_raw_run):
+    run = RunDirectory(mock_spb_raw_run)
+    count = run.get_data_counts('SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos.value')
+    assert count.index.tolist() == run.train_ids
+    assert (count.values == 1).all()