Fix labelled data counts not covering all selected train IDs

Author: philsmt

extra_data/keydata.py

@@ -398,7 +398,9 @@ def data_counts(self, labelled=True):
 
         if labelled:
             import pandas as pd
-            return pd.Series(counts, index=train_ids)
+            res = pd.Series(0, index=self.train_ids)
+            res.loc[train_ids] = counts
+            return res
         else:
             all_tids_arr = np.array(self.train_ids)
             res = np.zeros(len(all_tids_arr), dtype=np.uint64)

extra_data/tests/test_keydata.py

@@ -170,7 +170,7 @@ def test_get_train_keep_dims(mock_jungfrau_run):
     assert val.shape == (1, 16, 512, 1024)
 
 
-def test_data_counts(mock_reduced_spb_proc_run):
+def test_data_counts(mock_reduced_spb_proc_run, mock_jungfrau_run):
     run = RunDirectory(mock_reduced_spb_proc_run)
 
     # control data
@@ -189,6 +189,14 @@ def test_data_counts(mock_reduced_spb_proc_run):
     assert count.index.tolist() == mod.train_ids
     assert count.values.sum() == mod.shape[0]
 
+    # Combine data with more train IDs than our original sources, and
+    # ensure that the data counts now cover these additional train IDs
+    # as well with 0.
+    multi_run = run.union(RunDirectory(mock_jungfrau_run))
+    mod = multi_run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf', 'image.data']
+    count = mod.data_counts()
+    assert count.index.tolist() == mod.train_ids
+
 
 def test_data_counts_empty(mock_fxe_raw_run):
     run = RunDirectory(mock_fxe_raw_run)