Repository: European-XFEL/EXtra-data

Issue:
AssertionError in .data_counts(labelled=False)
@yohey-uemura discovered this over the weekend. When writing this method, I had overlooked the possibility of having a selected train which doesn't correspond to any data for the given source. For instance, in this run, LPD08 starts from train 1232492275, but other sources in the run start from 1232492274.

```python
run = open_run(2958, 172)
run[
    'FXE_DET_LPD1M-1/DET/8CH0:xtdf', 'image.cellId'
].data_counts(labelled=False)
```

```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-14-749f3505f4df> in <module>
      1 run[
      2     'FXE_DET_LPD1M-1/DET/8CH0:xtdf', 'image.cellId'
----> 3 ].data_counts(labelled=False)

/gpfs/exfel/sw/software/xfel_anaconda3/1.1.2/lib/python3.7/site-packages/extra_data/keydata.py in data_counts(self, labelled)
    168             # should be in order, but sometimes trains are written out of order.
    169             # Reorder the counts to match self.train_ids.
--> 170             assert len(train_ids) == len(self.train_ids)
    171             assert np.isin(train_ids, self.train_ids).all()
    172             idxs = np.argsort(train_ids)

AssertionError: 
```

Comments:

Patch:
diff --git a/.github/dependabot/constraints.txt b/.github/dependabot/constraints.txt
index ee3a594..7cfce9f 100644
--- a/.github/dependabot/constraints.txt
+++ b/.github/dependabot/constraints.txt
@@ -1,3 +1,4 @@
+coverage==6.1.2
 cycler==0.11.0
 h5py<3.7.0
 karabo-bridge==0.6.1
diff --git a/extra_data/keydata.py b/extra_data/keydata.py
index e81e88e..d753525 100644
--- a/extra_data/keydata.py
+++ b/extra_data/keydata.py
@@ -164,13 +164,16 @@ class KeyData:
             import pandas as pd
             return pd.Series(counts, index=train_ids)
         else:
-            # self.train_ids is always sorted. The train IDs from chunks
-            # should be in order, but sometimes trains are written out of order.
-            # Reorder the counts to match self.train_ids.
-            assert len(train_ids) == len(self.train_ids)
+            # We may be missing some train IDs, if they're not in any file
+            # for this source, and they're sometimes out of order within chunks
+            # (they shouldn't be, but we try not to fail too badly if they are).
             assert np.isin(train_ids, self.train_ids).all()
-            idxs = np.argsort(train_ids)
-            return counts[idxs]
+            tid_to_ix = {t: i for (i, t) in enumerate(self.train_ids)}
+            res = np.zeros(len(self.train_ids), dtype=np.uint64)
+            for tid, ct in zip(train_ids, counts):
+                res[tid_to_ix[tid]] = ct
+
+            return res
 
     def as_single_value(self, rtol=1e-5, atol=0.0, reduce_by='median'):
         """Retrieve a single reduced value if within tolerances.


Test patch:
diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py
index 2d107b1..3930cfb 100644
--- a/extra_data/tests/test_keydata.py
+++ b/extra_data/tests/test_keydata.py
@@ -2,8 +2,11 @@ import os
 import numpy as np
 import pytest
 
+import h5py
+
 from extra_data import RunDirectory, H5File
 from extra_data.exceptions import TrainIDError, NoDataError
+from . import make_examples
 from .mockdata import write_file
 from .mockdata.xgm import XGM
 
@@ -146,6 +149,46 @@ def test_data_counts_empty(mock_fxe_raw_run):
     count_none_arr = cam_nodata.drop_empty_trains().data_counts(labelled=False)
     assert len(count_none_arr) == 0
 
+
+@pytest.fixture()
+def fxe_run_module_offset(tmp_path):
+    run_dir = tmp_path / 'fxe-run-mod-offset'
+    run_dir.mkdir()
+    make_examples.make_fxe_run(run_dir, format_version='1.0')
+
+    # Shift the train IDs for a module by 1, so it has data for a different set
+    # of train IDs to other sources.
+    with h5py.File(run_dir / 'RAW-R0450-LPD08-S00000.h5', 'r+') as f:
+        tids_dset = f['INDEX/trainId']
+        tids_dset[:] = tids_dset[:] + 1
+
+    return run_dir
+
+
+def test_data_counts_missing_train(fxe_run_module_offset):
+    run = RunDirectory(fxe_run_module_offset)
+    assert len(run.train_ids) == 481
+    lpd_m8 = run['FXE_DET_LPD1M-1/DET/8CH0:xtdf', 'image.cellId']
+
+    ser = lpd_m8.data_counts(labelled=True)
+    assert len(ser) == 480
+    np.testing.assert_array_equal(ser.index, run.train_ids[1:])
+
+    arr = lpd_m8.data_counts(labelled=False)
+    assert len(arr) == 481
+    assert arr[0] == 0
+    np.testing.assert_array_equal(arr[1:], 128)
+
+    lpd_m8_w_data = lpd_m8.drop_empty_trains()
+    ser = lpd_m8_w_data.data_counts(labelled=True)
+    assert len(ser) == 480
+    np.testing.assert_array_equal(ser.index, run.train_ids[1:])
+
+    arr = lpd_m8_w_data.data_counts(labelled=False)
+    assert len(arr) == 480
+    np.testing.assert_array_equal(arr, 128)
+
+
 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']