Simplifies selectSASEinXGM() by using mask based on bunch pattern

xgm.py

@@ -183,7 +183,8 @@ def cleanXGMdata(data, npulses=None, sase3First=True):
 
 
 def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=None):
-    ''' Extract SASE1- or SASE3-only XGM data.
+    ''' Given an array containing both SASE1 and SASE3 data, extracts SASE1-
+        or SASE3-only XGM data.
         There are various cases depending on i) the mode of operation (10 Hz
         with fresh bunch, dedicated trains to one SASE, pulse on demand), 
         ii) the potential change of number of pulses per train in each SASE
@@ -192,7 +193,7 @@ def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=
         Inputs:
             data: xarray Dataset containing xgm data
             sase: key of sase to select: {'sase1', 'sase3'}
-            xgm: key of xgm to select: {'SA3_XGM', 'SCS_XGM'}
+            xgm: key of xgm to select: {'XTD10_XGM', 'SCS_XGM'}
             sase3First: bool, optional. Used in case no bunch pattern was recorded
             npulses: int, optional. Required in case no bunch pattern was recorded.
             
@@ -202,6 +203,7 @@ def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=
             that sase in the run. The missing values, in case of change of number of pulses,
             are filled with NaNs.
     '''
+    #1. case where bunch pattern is missing:
     if sase not in data:
         print('Missing bunch pattern info!')
         if npulses is None:
@@ -221,108 +223,27 @@ def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=
             else:
                 start=xgmData.shape[1]-npulses
                 return xgmData[:,start:start+npulses]
-    result = None
-    npulses_sa3 = data['npulses_sase3']       
-    npulses_sa1 = data['npulses_sase1']
-    dedicated = 0
-    if np.all(npulses_sa1.where(npulses_sa3 !=0, drop=True) == 0):
-        dedicated += 1
-        print('No SASE 1 pulses during SASE 3 operation')
-    if np.all(npulses_sa3.where(npulses_sa1 !=0, drop=True) == 0):
-        dedicated += 1
-        print('No SASE 3 pulses during SASE 1 operation')
-    #Alternating pattern with dedicated pulses in SASE1 and SASE3:
-    if dedicated==2:
-        if sase=='sase1':
-            result = data[xgm].where(npulses_sa1>0, drop=True)[:,:npulses_sa1.max().values]
-        else:
-            result = data[xgm].where(npulses_sa3>0, drop=True)[:,:npulses_sa3.max().values]
-        result = result.where(result != 1.0)
-        return result
-    # SASE1 and SASE3 bunches in a same train: find minimum indices of first and
-    # maximum indices of last pulse per train
-    else:
-        pulseIdmin_sa1 = data['sase1'].where(npulses_sa1 != 0).where(data['sase1']>1).min().values
-        pulseIdmax_sa1 = data['sase1'].where(npulses_sa1 != 0).where(data['sase1']>1).max().values
-        pulseIdmin_sa3 = data['sase3'].where(npulses_sa3 != 0).where(data['sase3']>1).min().values
-        pulseIdmax_sa3 = data['sase3'].where(npulses_sa3 != 0).where(data['sase3']>1).max().values
-        if pulseIdmin_sa1 > pulseIdmax_sa3:
-            sa3First = True
-        elif pulseIdmin_sa3 > pulseIdmax_sa1:
-            sa3First = False
-        else:
-            print('Interleaved mode, but no sase-dedicated XGM data loaded.')
-            saseStr = 'SA{}'.format(sase[4])
-            xgmStr = xgm.split('_')[0]
-            print('Loading {}_{} data...'.format(xgmStr, saseStr))
-            try:
-                if npulses == None:
-                    npulses = data['npulses_sase{}'.format(sase[4])].max().values
-                if xgmStr == 'XTD10':
-                    source = 'SA3_XTD10_XGM/XGM/DOOCS:output'
-                if xgmStr == 'SCS':
-                    source = 'SCS_BLU_XGM/XGM/DOOCS:output'
-                key = 'data.intensitySa{}TD'.format(sase[4])
-                result = data.attrs['run'].get_array(source, key, extra_dims=['XGMbunchId'])
-                result = result.isel(XGMbunchId=slice(0, npulses))                
-                return result
-            except:
-                print('Could not load {}_{} data. '.format(xgmStr, saseStr) +
-                  'Interleaved mode and no sase-dedicated data is not yet supported.')
-
-    #take the derivative along the trainId to track changes in pulse number:
-    diff = npulses_sa3.diff(dim='trainId')
-    #only keep trainIds where a change occured:
-    diff = diff.where(diff != 0, drop=True)
-    #get a list of indices where a change occured:
-    idx_change_sa3 = np.argwhere(np.isin(npulses_sa3.trainId.values,
-                                     diff.trainId.values, assume_unique=True))[:,0]
-
-    #Same for SASE 1:
-    diff = npulses_sa1.diff(dim='trainId')
-    diff = diff.where(diff !=0, drop=True)
-    idx_change_sa1 = np.argwhere(np.isin(npulses_sa1.trainId.values,
-                                     diff.trainId.values, assume_unique=True))[:,0]
-
-    #create index that locates all changes of pulse number in both SASE1 and 3:
-    #add index 0 to get the initial pulse number per train:
-    idx_change = np.unique(np.concatenate(([0], idx_change_sa3, idx_change_sa1))).astype(int)
+    
+    #2. case where bunch pattern is provided
+    xgm_arr = data[xgm].where(data[xgm] != 1., drop=True)
+    sa3 = data['sase3'].where(data['sase3'] > 1, drop=True)
+    sa3_val=np.unique(sa3)
+    sa3_val = sa3_val[~np.isnan(sa3_val)]
+    sa1 = data['sase1'].where(data['sase1'] > 1, drop=True)
+    sa1_val=np.unique(sa1)
+    sa1_val = sa1_val[~np.isnan(sa1_val)]
+    sa_all = xr.concat([sa1, sa3], dim='bunchId').rename('sa_all')
+    sa_all = xr.DataArray(np.sort(sa_all)[:,:xgm_arr['XGMbunchId'].shape[0]],
+                          dims=['trainId', 'bunchId'],
+                          coords={'trainId':data.trainId},
+                          name='sase_all')
+    mask_sa1 = sa_all.sel(trainId=sa1.trainId).isin(sa1_val).rename({'bunchId':'XGMbunchId'})
+    mask_sa3 = sa_all.sel(trainId=sa3.trainId).isin(sa3_val).rename({'bunchId':'XGMbunchId'})
     if sase=='sase1':
-        npulses = npulses_sa1
-        maxpulses = int(npulses_sa1.max().values)
+        return xgm_arr.where(mask_sa1, drop=True).rename('{}_SA1'.format(xgm.split('_')[0]))
     else:
-        npulses = npulses_sa3
-        maxpulses = int(npulses_sa3.max().values)
-    for i,k in enumerate(idx_change):    
-        #skip if no pulses after the change:
-        if npulses[idx_change[i]]==0:
-            continue
-        #calculate indices
-        if sa3First:
-            a = 0
-            b = int(npulses_sa3[k].values)
-            c = b
-            d = int(c + npulses_sa1[k].values)
-        else:
-            a = int(npulses_sa1[k].values)
-            b = int(a + npulses_sa3[k].values)
-            c = 0
-            d = a
-        if sase=='sase1':
-            a = c
-            b = d
-        if i==len(idx_change)-1:
-            l = None
-        else:
-            l = idx_change[i+1]
-        temp = data[xgm][k:l,a:a+maxpulses].copy()
-        temp[:,b:] = np.NaN
-        if result is None:
-            result = temp
-        else:
-            result = xr.concat([result, temp], dim='trainId')
-    return result
-
+        return xgm_arr.where(mask_sa3, drop=True).rename('{}_SA3'.format(xgm.split('_')[0]))
+    
 
 def saseContribution(data, sase='sase1', xgm='XTD10_XGM'):
     ''' Calculate the relative contribution of SASE 1 or SASE 3 pulses