Merge branch 'interleaved' into 'master'

Interleaved

See merge request !42

xgm.py

@@ -251,7 +251,24 @@ def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=
         elif pulseIdmin_sa3 > pulseIdmax_sa1:
             sa3First = False
         else:
-            print('Interleaved mode')
+            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')
@@ -578,6 +595,8 @@ def getTIMapd(data, mcp=1, use_apd=True, intstart=None, intstop=None,
     sa3 = data['sase3'].where(data['sase3']>1, drop=True)
     npulses_sa3 = data['npulses_sase3']
     maxpulses = int(npulses_sa3.max().values)
+    if npulses is not None:
+        maxpulses = np.min([npulses, maxpulses])
     step = 1
     if maxpulses > 1:
         #Calculate the number of non-lasing pulses between two lasing pulses (step)
@@ -599,7 +618,7 @@ def getTIMapd(data, mcp=1, use_apd=True, intstart=None, intstop=None,
     sa3 -= sa3[:,0]
     sa3 = sa3.astype(int)
     if np.all(npulses_sa3 == npulses_sa3[0]):
-        tim = apd[:, sa3[0].values]
+        tim = apd[:, sa3[0].values[:maxpulses]]
         return tim
     
     stride = 1
@@ -839,10 +858,11 @@ def checkTimApdWindow(data, mcp=1, use_apd=True, intstart=None, intstop=None):
     tid = data['npulses_sase3'].where(data['npulses_sase3'] == npulses_max,
                                       drop=True).trainId.values
     if 'MCP{}raw'.format(mcp) not in data:
-        ntid, data_from_train = data.attrs['run'].train_from_id(tid[0])
-        trace = data_from_train['SCS_UTC1_ADQ/ADC/1:network']['digitizers.channel_1_'
-                                +'{}.raw.samples'.format(mcpToChannel[mcp])]
-        print('no raw data for MCP{}. Loading trace from MCP{}'.format(mcp, mcp))
+        print('no raw data for MCP{}. Loading average trace from MCP{}'.format(mcp, mcp))
+        trace = data.attrs['run'].get_array(
+                'SCS_UTC1_ADQ/ADC/1:network',
+                'digitizers.channel_1_{}.raw.samples'.format(mcpToChannel[mcp])
+                ).sel({'trainId':tid}).mean(dim='trainId')
     else:
         trace = data['MCP{}raw'.format(mcp)].sel({'trainId':tid}).mean(dim='trainId')
     if use_apd: