Fixes for ADQ & pulse components with no trains of data

Author: takluyver

src/extra/components/adq.py

@@ -1137,7 +1137,10 @@ def pulse_data(self, labelled=True, pulse_dim='pulseId', train_roi=(),
         pulse_ids = pulses.pulse_ids(labelled=False)
 
         # Prepare output buffer shape.
-        out_shape = (len(pulse_ids), pulse_layout['length'].max())
+        if len(pulse_ids):
+            out_shape = (len(pulse_ids), pulse_layout['length'].max())
+        else:
+            out_shape = (0, 0)
 
         if parallel is not False:
             # Prepare parallelization.

src/extra/components/pulses.py

@@ -1145,7 +1145,10 @@ def _get_pulse_ids(self):
 
         if not counts:
             # Immediately return an empty series if there is no data.
-            return pd.Series([], dtype=np.int32)
+            index = pd.MultiIndex.from_arrays([
+                np.zeros(0, dtype=np.uint64), np.zeros(0, dtype=np.int64)
+            ], names=['trainId', 'pulseIndex'])
+            return pd.Series([], index=index, dtype=np.int32)
 
         index = pd.MultiIndex.from_arrays([
             np.repeat(self._key.train_id_coordinates(), counts),
@@ -1899,13 +1902,19 @@ def _get_pulse_ids(self):
             fel_by_train.append(np.isin(pids, fel_pids))
             ppl_by_train.append(np.isin(pids, ppl_pids))
 
+        def concat_1d(arrays, dtype):  # Allows an empty list of arrays
+            if len(arrays) == 0:
+                return np.zeros(0, dtype=dtype)
+            return np.concatenate(arrays)
+
         index = pd.MultiIndex.from_arrays([
             np.repeat(train_ids, counts),
-            np.concatenate([np.arange(count) for count in counts]),
-            np.concatenate(fel_by_train), np.concatenate(ppl_by_train)
+            concat_1d([np.arange(count) for count in counts], np.int64),
+            concat_1d(fel_by_train, np.bool_),
+            concat_1d(ppl_by_train, np.bool_),
         ], names=['trainId', 'pulseIndex', 'fel', 'ppl'])
 
-        return pd.Series(data=np.concatenate(pids_by_train),
+        return pd.Series(data=concat_1d(pids_by_train, np.int32),
                          index=index, dtype=np.int32)
 
     def _get_pulse_mask(self, reduced=False):

tests/test_components_adq.py

@@ -235,6 +235,11 @@ def test_pulse_data(mock_sqs_remi_run, parallel):
         # Test whether all pulses after the trace are nan.
         assert np.isnan(max_by_pulse[21:50]).all()
 
+    # No trains in data
+    empty_data = ch.select_trains(np.s_[:0]).pulse_data()
+    assert empty_data.dims == ('pulse', 'sample')
+    assert empty_data.shape[0] == 0
+
     # Use pulse information with not all trains available.
     ch = AdqRawChannel(
         mock_sqs_remi_run, '3B', digitizer='SQS_DIGITIZER_UTC2',

tests/test_components_pulses.py

@@ -364,6 +364,11 @@ def test_pulse_ids(mock_spb_aux_run, source):
     # Test unlabelled.
     np.testing.assert_equal(pulses.pulse_ids(labelled=False), pids)
 
+    # Test with no trains selected
+    pids_empty = pulses.select_trains(np.s_[:0]).pulse_ids()
+    assert len(pids_empty) == 0
+    assert pids_empty.index.names == ['trainId', 'pulseIndex']
+
     # Test deprecated method.
     with pytest.warns():
         assert pulses.get_pulse_ids().equals(pulses.pulse_ids())
@@ -670,6 +675,11 @@ def test_pump_probe_basic(mock_spb_aux_run, source):
     assert not ppl[0]
     assert ppl[1:51].all()
 
+    # Pulse IDs with no data
+    pids_empty = pulses.select_trains(np.s_[:0]).pulse_ids()
+    assert len(pids_empty) == 0
+    assert pids_empty.index.names == ['trainId', 'pulseIndex', 'fel', 'ppl']
+
     # Pulse mask.
     assert pulses.pulse_mask(labelled=False)[0, 1000:1306:6].all()