Adds XGM calibration function

xgm.py

@@ -187,6 +187,7 @@ def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM'):
             result = xr.concat([result, temp], dim='trainId')
     return result
 
+
 def calcContribSASE(data, sase='sase1', xgm='SA3_XGM'):
     ''' Calculate the relative contribution of SASE 1 or SASE 3 pulses 
         for each train in the run. Supports fresh bunch, dedicated trains
@@ -216,6 +217,7 @@ def calcContribSASE(data, sase='sase1', xgm='SA3_XGM'):
     else:
         return 1 - contrib
 
+
 def filterOnTrains(data, key='sase3'):
     ''' Removes train ids for which there was no pulse in sase='sase1' or 'sase3' branch
         
@@ -230,6 +232,110 @@ def filterOnTrains(data, key='sase3'):
     res = data.where(data[key]>0, drop=True)
     return res
 
+
+def calibrateXGMs(data, rollingWindow=200, plot=False):
+    ''' Calibrate the fast (pulse-resolved) signals of the XTD10 and SCS XGM 
+        (read in intensityTD property) to the respective slow ion signal 
+        (photocurrent read by Keithley, channel 'pulseEnergy.photonFlux.value').
+        One has to take into account the possible signal created by SASE1 pulses. In the
+        tunnel, this signal is usually large enough to be read by the XGM and the relative
+        contribution C of SASE3 pulses to the overall signal is computed.
+        In the tunnel, the calibration F is defined as:
+            F = E_slow / E_fast_avg, where
+        E_fast_avg is the rolling average (with window rollingWindow) of the fast signal.
+        In SCS XGM, the signal from SASE1 is usually in the noise, so we calculate the 
+        average over the pulse-resolved signal of SASE3 pulses only and calibrate it to the 
+        slow signal modulated by the SASE3 contribution:
+            F = (N1+N3) * E_avg * C/(N3 * E_fast_avg_sase3), where N1 and N3 are the number 
+        of pulses in SASE1 and SASE3, E_fast_avg_sase3 is the rolling average (with window
+        rollingWindow) of the SASE3-only fast signal.
+        
+        Inputs:
+            data: xarray Dataset
+            rollingWindow: length of running average to calculate E_fast_avg
+            plot: boolean, plot the calibration output
+            
+        Output:
+            factors: numpy ndarray of shape 1 x 2 containing 
+                     [XTD10 calibration factor, SCS calibration factor]
+    '''
+    noSCS = noSA3 = False
+    sa3_calib_factor = None
+    scs_calib_factor = None
+    if 'SCS_XGM' not in data:
+        print('no SCS XGM data. Skipping calibration for SCS XGM')
+        noSCS = True
+    if 'SA3_XGM' not in data:
+        print('no SASE3 XGM data. Skipping calibration for SASE3 XGM')
+        noSA3 = True
+    if noSCS and noSA3:
+        return np.array([None, None])
+    start = 0
+    stop = None
+    npulses = data['npulses_sase3']
+    ntrains = npulses.shape[0]
+    # First, in case of change in number of pulses, locate a region where
+    # the number of pulses is maximum.
+    if not np.all(npulses == npulses[0]):
+        print('Warning: Number of pulses per train changed during the run!')
+        start = np.argmax(npulses.values)
+        stop = ntrains + np.argmax(npulses.values[::-1]) - 1
+        if stop - start < rollingWindow:
+            print('not enough consecutive data points with the largest number of pulses per train')
+        start += rollingWindow
+        stop = np.min((ntrains, stop+rollingWindow))
+    # Calibrate SASE3 XGM with all signal from SASE1 and SASE3
+    if not noSA3:
+        xgm_avg = data['SA3_XGM'].where(data['SA3_XGM'] != 1.0).mean(axis=1)
+        rolling_sa3_xgm = xgm_avg.rolling(trainId=rollingWindow).mean()
+        ratio = data['SA3_XGM_SLOW']/rolling_sa3_xgm
+        sa3_calib_factor = ratio[start:stop].mean().values
+        print('calibration factor SA3 XGM: %f'%sa3_calib_factor)
+
+    # Calibrate SCS XGM with SASE3-only contribution
+    sa3contrib = calcContribSASE(data, 'sase3', 'SA3_XGM')
+    if not noSCS:
+        scs_sase3_fast = selectSASEinXGM(data, 'sase3', 'SCS_XGM').mean(axis=1)
+        meanFast = scs_sase3_fast.rolling(trainId=rollingWindow).mean()
+        ratio = ((data['npulses_sase3']+data['npulses_sase1']) *
+                 data['SCS_XGM_SLOW'] * sa3contrib) / (meanFast * data['npulses_sase3'])
+        scs_calib_factor = ratio[start:stop].median().values
+        print('calibration factor SCS XGM: %f'%scs_calib_factor)
+        
+    if plot:
+        plt.figure(figsize=(8,8))
+        plt.subplot(211)
+        plt.title('E[uJ] = %.2f x IntensityTD' %(sa3_calib_factor))
+        plt.plot(data['SA3_XGM_SLOW'], label='SA3 slow', color='C1')
+        plt.plot(rolling_sa3_xgm*sa3_calib_factor,
+                 label='SA3 fast signal rolling avg', color='C4')
+        plt.ylabel('Energy [uJ]')
+        plt.xlabel('train in run')
+        plt.legend(loc='upper left', fontsize=10)
+        plt.twinx()
+        plt.plot(xgm_avg*sa3_calib_factor, label='SA3 fast signal train avg', alpha=0.2, color='C4')
+        plt.ylabel('Calibrated SA3 fast signal [uJ]')
+        plt.legend(loc='lower right', fontsize=10)
+
+        plt.subplot(212)
+        plt.title('E[uJ] = %.2f x HAMP' %scs_calib_factor)
+        plt.plot(data['SCS_XGM_SLOW'], label='SCS slow (all SASE)', color='C0')
+        slow_avg_sase3 = data['SCS_XGM_SLOW']*(data['npulses_sase1']
+                                                    +data['npulses_sase3'])*sa3contrib/data['npulses_sase3']
+        plt.plot(slow_avg_sase3, label='SCS slow (SASE3 only)', color='C1')
+        plt.plot(meanFast*scs_calib_factor, label='SCS HAMP rolling avg', color='C2')
+        plt.ylabel('Energy [uJ]')
+        plt.xlabel('train in run')
+        plt.legend(loc='upper left', fontsize=10)
+        plt.twinx()
+        plt.plot(scs_sase3_fast*scs_calib_factor, label='SCS HAMP train avg', alpha=0.2, color='C2')
+        plt.ylabel('Calibrated HAMP signal [uJ]')
+        plt.legend(loc='lower right', fontsize=10)
+        plt.tight_layout()
+
+    return np.array([sa3_calib_factor, scs_calib_factor])
+
+
 def mcpPeaks(data, intstart, intstop, bkgstart, bkgstop, t_offset=1760, mcp=1, npulses=None):
     ''' Computes peak integration from raw MCP traces.
     
@@ -265,6 +371,7 @@ def mcpPeaks(data, intstart, intstop, bkgstart, bkgstop, t_offset=1760, mcp=1, n
         results[:,i] = np.trapz(data[keyraw][:,a:b] - bg, axis=1)
     return results
 
+
 def getTIMapd(data, mcp=1, use_apd=True, intstart=None, intstop=None,
               bkgstart=None, bkgstop=None, t_offset=1760, npulses=None):
     ''' Extract peak-integrated data from TIM where pulses are from SASE3 only.
@@ -323,6 +430,7 @@ def getTIMapd(data, mcp=1, use_apd=True, intstart=None, intstop=None,
             tim = xr.concat([tim, temp], dim='trainId')
     return tim
     
+
 def calibrateTIM(data, rollingWindow=200, mcp=1, use_apd=True, intstart=None, intstop=None,
               bkgstart=None, bkgstop=None, t_offset=1760, npulses_apd=None):
     ''' Calibrate TIM signal (Peak-integrated signal) to the slow ion signal of SCS_XGM
@@ -343,8 +451,7 @@ def calibrateTIM(data, rollingWindow=200, mcp=1, use_apd=True, intstart=None, in
         
         Inputs:
             data: xarray Dataset
-            rolling window: number of trains to perform a running average on to match
-                            TIM-avg and E_slow
+            rollingWindow: length of running average to calculate TIM_avg
             mcp: MCP channel
             use_apd: boolean. If False, the TIM pulse peaks are extract from raw traces using
                      getTIMapd
@@ -377,5 +484,6 @@ def calibrateTIM(data, rollingWindow=200, mcp=1, use_apd=True, intstart=None, in
     ratio = ((data['npulses_sase3']+data['npulses_sase1']) *
              data['SCS_XGM_SLOW'] * sa3contrib) / (avgFast*data['npulses_sase3'])
     F = float(ratio[start:stop].median().values)
+    #print('calibration factor TIM: %f'%F)
     return F