diff --git a/src/toolbox_scs/detectors/pes.py b/src/toolbox_scs/detectors/pes.py index ff3a8220c15ec0d64802c10eeaf8f7a4e6897153..e1bc70b992eb5455f9c632b99ebe315586bd00b9 100644 --- a/src/toolbox_scs/detectors/pes.py +++ b/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