diff --git a/bunch_pattern.py b/bunch_pattern.py index 42f34a61f8724762be1c3670d1287b74a64ac5d6..b74f5344dccce706d48da1222e6fe1fc4385bf3b 100644 --- a/bunch_pattern.py +++ b/bunch_pattern.py @@ -9,6 +9,9 @@ import numpy as np import xarray as xr import ToolBox as tb +import os +from extra_data.read_machinery import find_proposal +from extra_data import RunDirectory def extractBunchPattern(bp_table=None, key='sase3', runDir=None): ''' generate the bunch pattern and number of pulses of a source directly from the @@ -161,24 +164,49 @@ def pulsePatternInfo(data, plot=False): plt.tight_layout() -def repRate(data, sase='sase3'): +def repRate(data=None, runNB=None, proposalNB=None, key='sase3'): ''' Calculates the pulse repetition rate (in kHz) in sase according to the bunch pattern and assuming a grid of 4.5 MHz. Inputs: - data: xarray Dataset containing pulse pattern - sase: sase in which the repetition rate is - calculated (1,2 or 3) + data: xarray Dataset containing pulse pattern, needed if runNB is none + runNB: int or str, run number. Needed if data is None + proposal: int or str, proposal where to find the run. Needed if data is None + key: str in [sase1, sase2, sase3, scs_ppl], source for which the + repetition rate is calculated Output: f: repetition rate in kHz ''' - assert sase in data, 'key "{}" not found in data!'.format(sase) - sase = data[sase].where(data['npulses_{}'.format(sase)]>1, - drop=True).values - if len(sase)==0: + if runNB is None and data is None: + raise ValueError('Please provide either the runNB + proposal or the data argument.') + if runNB is not None and proposalNB is None: + raise ValueError('Proposal is missing.') + if runNB is not None: + if isinstance(runNB, int): + runNB = 'r{:04d}'.format(runNB) + if isinstance(proposalNB,int): + proposalNB = 'p{:06d}'.format(proposalNB) + runFolder = os.path.join(find_proposal(proposalNB), 'raw', runNB) + runDir = RunDirectory(runFolder) + bp_mnemo = tb.mnemonics['bunchPatternTable'] + if bp_mnemo['source'] not in runDir.all_sources: + raise ValueError('Source {} not found in run'.format( + bp_mnemo['source'])) + else: + bp_table = runDir.get_array(bp_mnemo['source'],bp_mnemo['key'], + extra_dims=bp_mnemo['dim']) + a, b, mask = extractBunchPattern(bp_table, key=key) + else: + if key not in ['sase1', 'sase3']: + a, b, mask = extractBunchPattern(key=key, runDir=data.attrs['run']) + else: + a = data[key] + b = data[f'npulses_{key}'] + a = a.where(b > 1, drop = True).values + if len(a)==0: print('Not enough pulses to extract repetition rate') return 0 - f = 1/((sase[0,1] - sase[0,0])*12e-3/54.1666667) + f = 1/((a[0,1] - a[0,0])*12e-3/54.1666667) return f def sortBAMdata(data, key='sase3'):