PES: Check the size of raw trace against pulse pattern

src/toolbox_scs/detectors/pes.py

@@ -27,6 +27,7 @@ log = logging.getLogger(__name__)
 
 def get_pes_tof(run, mnemonics=None, merge_with=None,
                 start=31390, width=300, origin=None, width_ns=None,
+                subtract_baseline=True,
                 baseStart=None, baseWidth=80,
                 sample_rate=2e9):
     """
@@ -57,6 +58,9 @@ def get_pes_tof(run, mnemonics=None, merge_with=None,
     width_ns: float
         time window for one spectrum. If None, the time window is defined by
         width / sample rate.
+    subtract_baseline: bool
+        If True, subtract baseline defined by baseStart and baseWidth to each
+        spectrum.
     baseStart: int
         starting sample of the baseline.
     baseWidth: int
@@ -114,26 +118,35 @@ def get_pes_tof(run, mnemonics=None, merge_with=None,
             arr = merge_with[m]
         else:
             arr = run.get_array(*run_mnemonics[m].values(), name=m)
+        if arr.sizes['PESsampleId'] < npulses*period*440 + start + width:
+            log.warning('Not all pulses were recorded. The number of samples on '
+                        f'the digitizer {arr.sizes["PESsampleId"]} is not enough '
+                        f'to cover the {npulses} spectra. Missing spectra will be '
+                       'filled with NaNs.')
         spectra = []
         for p in range(npulses):
             begin = p*period*440 + start
             end = begin + width
-            baseBegin = p*period*440 + baseStart
-            baseEnd = baseBegin + baseWidth
+            if end > arr.sizes['PESsampleId']:
+                break
             pes = arr.isel(PESsampleId=slice(begin, end))
-            bl = arr.isel(
-                PESsampleId=slice(baseBegin, baseEnd)).mean(dim='PESsampleId')
-            spectra.append(pes - bl)
+            if subtract_baseline:
+                baseBegin = p*period*440 + baseStart
+                baseEnd = baseBegin + baseWidth
+                bl = arr.isel(
+                    PESsampleId=slice(baseBegin, baseEnd)).mean(dim='PESsampleId')
+                pes = pes - bl
+            spectra.append(pes)
         spectra = xr.concat(spectra,
                             dim='sa3_pId').rename(m.replace('raw', 'tof'))
         ds = ds.merge(spectra)
     if len(ds.variables) > 0:
-        ds = ds.assign_coords({'sa3_pId': mask_on['pulse_slot'].values})
+        ds = ds.assign_coords({'sa3_pId': mask_on['pulse_slot'][:ds.sizes['sa3_pId']].values})
         ds = ds.rename({'PESsampleId': 'time_ns'})
         ds = ds.assign_coords({'time_ns': time_ns})
     if bool(merge_with):
         ds = merge_with.drop(to_process,
-                             errors='ignore').merge(ds, join='inner')
+                             errors='ignore').merge(ds, join='left')
 
     return ds