Improved repRate() function

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'):