diff --git a/docs/changelog.md b/docs/changelog.md index 79f30a62..fb613026 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -45,6 +45,8 @@ Fixed: [DataArray][xarray.DataArray] inputs as well (!419). - Fix incorrect time units when using [OpticalLaserDelay][extra.components.OpticalLaserDelay] with pre-2022 BAM data (!425). +- [PulsePattern][extra.components.pulses.PulsePattern]-based components now properly account for + all selected trains internally, not only those with data records (!434). Changed: diff --git a/src/extra/components/pulses.py b/src/extra/components/pulses.py index 0bacb958..8eda5e50 100644 --- a/src/extra/components/pulses.py +++ b/src/extra/components/pulses.py @@ -61,16 +61,21 @@ def _get_train_ids(self): This method may be overriden by any implementation of this class for either performance or if the underlying data may contain train IDs that have no pulse IDs associated with it. The default - implementation draws train IDs from the series of pulse IDs, and - thus cannot contain trains without pulses. This is particularly - relevant when counting pulses. + implementation draws train IDs from the underlying source or key + if available. Otherwise, the trains occuring in the series of + pulse IDs is used, which is therefore not able to include trains + without pulses. Returns: (np.ndarray) Train IDs, expected to be in order. """ - # This method turned out to be the fastest to get just the - # group labels. + if self._source is not None: + return self._source.train_ids + elif self._key is not None: + return self._key.train_ids + + # Fallback using pulse IDs. return self._get_pulse_ids().index.to_frame()['trainId'].unique() def _get_pulse_ids(self): @@ -783,9 +788,6 @@ def _find_pulsepattern_source(cls, data): raise ValueError('no timeserver or ppdecoder found, please pass ' 'one explicitly') - def _get_train_ids(self): - return self._key.train_id_coordinates() - def _get_pulse_ids(self): if self._with_timeserver: pids_by_train = [np.flatnonzero(mask) for mask @@ -1548,9 +1550,6 @@ def __init__(self, detector, *, clock_ratio=None, first_pulse_id=None, self._first_pulse_id = first_pulse_id self._negative_ppl_indices = negative_ppl_indices - def _get_train_ids(self): - return np.unique(self._key.train_id_coordinates()) - def _get_pulse_ids(self): triggers = self._key.ndarray() train_ids = self._key.train_id_coordinates() diff --git a/tests/test_components_pulses.py b/tests/test_components_pulses.py index 49938b25..5acac525 100644 --- a/tests/test_components_pulses.py +++ b/tests/test_components_pulses.py @@ -99,10 +99,14 @@ def test_init(mock_spb_aux_run): ids=['timeserver-control', 'timeserver-instrument', 'ppdecoder'] ) def test_no_trains(mock_spb_aux_run, source): - # Test with entirely empty data. pulses = XrayPulses(mock_spb_aux_run, source) + + # No pulses at all in this data. assert pulses.pulse_ids().empty - assert pulses.pulse_counts().empty + + # But trains (without pulses). + assert len(pulses.pulse_counts()) == 100 + np.testing.assert_array_equal(pulses.pulse_counts(), 0) def test_select_trains(mock_spb_aux_run): @@ -563,11 +567,13 @@ def test_dld_pulses(capsys): triggers['ppl'] = False mock_key = MagicMock() + mock_key.train_ids = [1000] mock_key.train_id_coordinates.return_value = np.repeat(1000, 10) mock_key.data_counts.return_value = np.array([10]) mock_key.ndarray.return_value = triggers mock_source = MagicMock() + mock_source.train_ids = [1000] mock_source.__getitem__ = lambda self, _: mock_key # Test regular.