diff --git a/src/toolbox_scs/detectors/digitizers.py b/src/toolbox_scs/detectors/digitizers.py index 26c06cfff116a6ca71d3af8a544d137ce33e76cd..a3ca5796c04253e1630f24d29f7ff9c5b24a544e 100644 --- a/src/toolbox_scs/detectors/digitizers.py +++ b/src/toolbox_scs/detectors/digitizers.py @@ -914,6 +914,9 @@ def get_digitizer_peaks(run, mnemonics, digitizer, merge_with=None, elif 'bunchPatternTable' in run_mnemonics: bpt = run.get_array(*run_mnemonics['bunchPatternTable'].values()) log.debug('Loaded bpt from DataCollection.') + elif 'bunchPatternTable_SA3' in run_mnemonics: + bpt = run.get_array(*run_mnemonics['bunchPatternTable_SA3'].values()) + log.debug('Loaded bpt from DataCollection.') else: bpt = None diff --git a/src/toolbox_scs/load.py b/src/toolbox_scs/load.py index 48360c79ec152c8221a2d705949d13d20743e3be..52b969776ec62490766d420da40863183d41e66b 100644 --- a/src/toolbox_scs/load.py +++ b/src/toolbox_scs/load.py @@ -123,20 +123,13 @@ def load(proposalNB=None, runNB=None, run.info() data_arrays = [] - run_mnemonics = mnemonics_for_run(run) # load pulse pattern info - if 'bunchPatternTable' in run_mnemonics: - bpt = run.get_array(*run_mnemonics['bunchPatternTable'].values(), - name='bunchPatternTable') - data_arrays.append(bpt) - elif 'bunchPatternTable_SA3' in run_mnemonics: - log.info('Did not find SCS bunch pattern table but found the SA3 one.') - bpt = run.get_array(*run_mnemonics['bunchPatternTable_SA3'].values(), - name='bunchPatternTable') - data_arrays.append(bpt) - else: + bpt = load_bpt(run, run_mnemonics=run_mnemonics) + if bpt is None: log.warning('Bunch pattern table not found in run. Skipping!') + else: + data_arrays.append(bpt) for f in fields: if type(f) == dict: @@ -322,3 +315,20 @@ def get_array(run, mnemonic_key=None, stepsize=None): raise return data + + +def load_bpt(run, merge_with=None, run_mnemonics=None): + if run_mnemonics is None: + run_mnemonics = mnemonics_for_run(run) + + for key in ['bunchPatternTable', 'bunchPatternTable_SA3']: + if bool(merge_with) and key in merge_with: + log.debug(f'Using {key} from merge_with dataset.') + return merge_with[key] + if key in run_mnemonics: + bpt = run.get_array(*run_mnemonics[key].values(), + name='bunchPatternTable') + log.debug(f'Loaded {key} from DataCollection.') + return bpt + log.debug('Could not find bunch pattern table.') + return None