Always use all selected trains in pulse pattern components

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:
 

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()

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.