diff --git a/extra_data/sourcedata.py b/extra_data/sourcedata.py index cf361511..cd358cd6 100644 --- a/extra_data/sourcedata.py +++ b/extra_data/sourcedata.py @@ -405,6 +405,13 @@ def data_counts(self, labelled=True, index_group=None): IDs. Otherwise, returns a NumPy array of counts to match ``.train_ids``. """ + def build_data_counts(data_counts): + if labelled: + import pandas as pd + return pd.DataFrame(data_counts).max(axis=1) + else: + return np.stack(list(data_counts.values())).max(axis=0) + if index_group is None: # Collect data counts for a sample key per index group. data_counts = { @@ -416,15 +423,17 @@ def data_counts(self, labelled=True, index_group=None): if not data_counts: data_counts = {None: np.zeros(len(self.train_ids), dtype=int)} - if labelled: - import pandas as pd - return pd.DataFrame(data_counts).max(axis=1) - else: - return np.stack(list(data_counts.values())).max(axis=0) + return build_data_counts(data_counts) else: - return self[self.one_key(index_group)].data_counts( - labelled=labelled) + if (key := self.one_key(index_group)) is None: + # Index group is actually keyless and not ownly + # downselected, most likely to occur for RUN-only + # sources. + return build_data_counts( + {None: np.zeros(len(self.train_ids), dtype=int)}) + + return self[key].data_counts(labelled=labelled) def train_id_coordinates(self, index_group=None): """Make an array of train IDs to use alongside data this source. diff --git a/extra_data/tests/test_sourcedata.py b/extra_data/tests/test_sourcedata.py index fc73064d..d021c449 100644 --- a/extra_data/tests/test_sourcedata.py +++ b/extra_data/tests/test_sourcedata.py @@ -360,3 +360,4 @@ def test_no_control_keys(mock_remi_run): assert sd.one_key() is None assert sd.aggregator == 'REMI01' np.testing.assert_array_equal(sd.data_counts(), 0) + np.testing.assert_array_equal(sd.data_counts(''), 0)