From 86303ce34ba7552bc20d22aa9a2d9b013caffc12 Mon Sep 17 00:00:00 2001 From: Philipp Schmidt Date: Thu, 2 Jul 2026 08:35:18 +0200 Subject: [PATCH] Fix labelled data counts not covering all selected train IDs --- extra_data/keydata.py | 4 +++- extra_data/tests/test_keydata.py | 14 +++++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/extra_data/keydata.py b/extra_data/keydata.py index 94bcf87c..02c83f5c 100644 --- a/extra_data/keydata.py +++ b/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) diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py index 52727727..9f3acc0b 100644 --- a/extra_data/tests/test_keydata.py +++ b/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) @@ -230,8 +238,8 @@ def test_data_counts_missing_train(fxe_run_module_offset): 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:]) + assert len(ser) == 481 + np.testing.assert_array_equal(ser.index, run.train_ids) arr = lpd_m8.data_counts(labelled=False) assert len(arr) == 481