diff --git a/Load.py b/Load.py index 6ed697b59e8b0960149a52ce7e90157e9a433da4..a571d3fac12f5ab561393070cf997e2047884766 100644 --- a/Load.py +++ b/Load.py @@ -445,8 +445,8 @@ def load(fields, runNB, proposalNB, subFolder='raw', display=False, validate=Fal else: bp_table = run.get_array(bp_mnemo['source'],bp_mnemo['key'], extra_dims=bp_mnemo['dim']) - sase1, npulses_sase1 = extractBunchPattern(bp_table, 'sase1') - sase3, npulses_sase3 = extractBunchPattern(bp_table, 'sase3') + sase1, npulses_sase1, dummy = extractBunchPattern(bp_table, 'sase1') + sase3, npulses_sase3, dummy = extractBunchPattern(bp_table, 'sase3') keys += ["sase1", "npulses_sase1", "sase3", "npulses_sase3"] vals += [sase1, npulses_sase1, sase3, npulses_sase3] else: diff --git a/bunch_pattern.py b/bunch_pattern.py index 8e2b8557a8a02bf412a5f12f6a0dec48b707341b..f51e60f50276e2fbae9fb047e5be689fa35b2290 100644 --- a/bunch_pattern.py +++ b/bunch_pattern.py @@ -25,6 +25,7 @@ def extractBunchPattern(bp_table=None, key='sase3', runDir=None): bunchPattern: DataArray containing indices of the sase/laser pulses for each train npulses: DataArray containing the number of pulses for each train + matched: 2-D DataArray mask (trainId x 2700), True where 'key' has pulses ''' keys=['sase1', 'sase2', 'sase3', 'scs_ppl'] @@ -74,7 +75,7 @@ def extractBunchPattern(bp_table=None, key='sase3', runDir=None): npulses = xr.DataArray(npulses, dims=['trainId'], coords={'trainId':matched.trainId}, name=f'npulses_{key}') - return bunchPattern, npulses + return bunchPattern, npulses, matched def pulsePatternInfo(data, plot=False): @@ -179,3 +180,30 @@ def repRate(data, sase='sase3'): return 0 f = 1/((sase[0,1] - sase[0,0])*12e-3/54.1666667) return f + +def sortBAMdata(data, key='sase3'): + ''' Extracts beam arrival monitor data from the raw arrays 'BAM6', 'BAM7', etc... + according to the bunchPatternTable. + Inputs: + data: xarray Dataset containing BAM arrays + key: str, ['sase1', 'sase2', 'sase3', 'scs_ppl'] + + Output: + ndata: xarray Dataset with same keys as input data (but new bam arrays) + ''' + a, b, mask = extractBunchPattern(key=key, runDir=data.attrs['run']) + mask = mask.rename({'pulse_slot':'BAMbunchId'}) + ndata = data + dropList = [] + mergeList = [] + for key in data: + if 'BAM' in key: + dropList.append(key) + bam = data[key].isel(BAMbunchId=slice(0,5400,2)) + bam = bam.where(mask, drop=True) + mergeList.append(bam) + mergeList.append(data.drop(dropList)) + ndata = xr.merge(mergeList, join='inner') + for k in data.attrs.keys(): + ndata.attrs[k] = data.attrs[k] + return ndata \ No newline at end of file