diff --git a/bunch_pattern.py b/bunch_pattern.py index b135db4575949bde132210218a421044312950fd..42f34a61f8724762be1c3670d1287b74a64ac5d6 100644 --- a/bunch_pattern.py +++ b/bunch_pattern.py @@ -187,7 +187,8 @@ def sortBAMdata(data, key='sase3'): corresponds to FLASH busrt length of 800 us @ 9 MHz. The bunchPatternTable only has 2700 values, corresponding to XFEL 600 us burst length @ 4.5 MHz. Hence, we truncate the BAM arrays to 5400 with a stride of 2 and match them - to the bunchPatternTable. + to the bunchPatternTable. If key is one of the sase, the given dimension name + of the bam arrays is 'sa[sase number]_pId', to match other data (XGM, TIM...). Inputs: data: xarray Dataset containing BAM arrays key: str, ['sase1', 'sase2', 'sase3', 'scs_ppl'] @@ -200,11 +201,14 @@ def sortBAMdata(data, key='sase3'): ndata = data dropList = [] mergeList = [] - for key in data: - if 'BAM' in key: - dropList.append(key) - bam = data[key].isel(BAMbunchId=slice(0,5400,2)) + for k in data: + if 'BAM' in k: + dropList.append(k) + bam = data[k].isel(BAMbunchId=slice(0,5400,2)) bam = bam.where(mask, drop=True) + if 'sase' in key: + name = f'sa{key[4]}_pId' + bam = bam.rename({'BAMbunchId':name}) mergeList.append(bam) mergeList.append(data.drop(dropList)) ndata = xr.merge(mergeList, join='inner')