Add plotting option to calibrateTIM()

Add RunDirectory as attribute of the DataSet in load()

Load.py

-# -*- coding: utf-8 -*-
-""" Toolbox for SCS.
-
-    Various utilities function to quickly process data measured at the SCS instruments.
-
-    Copyright (2019) SCS Team.
-"""
-
 import numpy as np
 from karabo_data import RunDirectory
 import xarray as xr
@@ -221,4 +213,6 @@ def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, vali
 
     aligned_vals = xr.align(*vals, join='inner')
     result = dict(zip(keys, aligned_vals))
-    return xr.Dataset(result)
+    result = xr.Dataset(result)
+    result.attrs['run'] = run
+    return result

xgm.py

+# -*- coding: utf-8 -*-
+""" Toolbox for SCS.
+
+    Various utilities function to quickly process data measured at the SCS instruments.
+
+    Copyright (2019) SCS Team.
+"""
 import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
@@ -429,9 +436,9 @@ def getTIMapd(data, mcp=1, use_apd=True, intstart=None, intstop=None,
         else:
             tim = xr.concat([tim, temp], dim='trainId')
     return tim
-    
 
-def calibrateTIM(data, rollingWindow=200, mcp=1, use_apd=True, intstart=None, intstop=None,
+
+def calibrateTIM(data, rollingWindow=200, mcp=1, plot=False, 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
         (photocurrent read by Keithley, channel 'pulseEnergy.photonFlux.value').
@@ -453,6 +460,7 @@ def calibrateTIM(data, rollingWindow=200, mcp=1, use_apd=True, intstart=None, in
             data: xarray Dataset
             rollingWindow: length of running average to calculate TIM_avg
             mcp: MCP channel
+            plot: boolean. If True, plot calibration results.
             use_apd: boolean. If False, the TIM pulse peaks are extract from raw traces using
                      getTIMapd
             intstart: trace index of integration start
@@ -484,6 +492,75 @@ 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
 
+    if plot:
+        fig = plt.figure(figsize=(8,5))
+        ax = plt.subplot(211)
+        ax.set_title('E[uJ] = {:2e} x TIM (MCP{})'.format(F, mcp))
+        ax.plot(data['SCS_XGM_SLOW'], label='SCS XGM slow (all SASE)', color='C0')
+        slow_avg_sase3 = data['SCS_XGM_SLOW']*(data['npulses_sase1']
+                                                    +data['npulses_sase3'])*sa3contrib/data['npulses_sase3']
+        ax.plot(slow_avg_sase3, label='SCS XGM slow (SASE3 only)', color='C1')
+        ax.plot(avgFast*F, label='Calibrated TIM rolling avg', color='C2')
+        ax.legend(loc='upper left', fontsize=8)
+        ax.set_ylabel('Energy [$\mu$J]', size=10)
+        #ax2=ax#.twinx()
+        ax.plot(filteredTIM.mean(axis=1)*F, label='Calibrated TIM train avg', alpha=0.2, color='C9')
+        #ax2.set_ylabel('Calibrated TIM (MCP{}) [uJ]'.format(mcp))
+        ax.legend(loc='best', fontsize=8, ncol=2)
+        plt.xlabel('train in run')
+        
+        ax = plt.subplot(234)
+        xgm_fast = selectSASEinXGM(data)
+        ax.scatter(filteredTIM, xgm_fast, s=5, alpha=0.1)
+        fit, cov = np.polyfit(filteredTIM.values.flatten(),xgm_fast.values.flatten(),1, cov=True)
+        y=np.poly1d(fit)
+        x=np.linspace(filteredTIM.min(), filteredTIM.max(), 10)
+        ax.plot(x, y(x), lw=2, color='r')
+        ax.set_ylabel('Raw HAMP [$\mu$J]', size=10)
+        ax.set_xlabel('TIM (MCP{}) signal'.format(mcp), size=10)
+        ax.annotate(s='y(x) = F x + A\n'+
+                    'F = %.3e\n$\Delta$F/F = %.2e\n'%(fit[0],np.abs(np.sqrt(cov[0,0])/fit[0]))+
+                    'A = %.3e'%fit[1],
+                    xy=(0.5,0.6), xycoords='axes fraction', fontsize=10, color='r')
+        print('TIM calibration factor: %e'%(F))
+        
+        ax = plt.subplot(235)
+        ax.hist(filteredTIM.values.flatten()*F, bins=50, rwidth=0.8)
+        ax.set_ylabel('number of pulses', size=10)
+        ax.set_xlabel('Pulse energy MCP{} [uJ]'.format(mcp), size=10)
+        #ax2 = ax.twiny()
+        #ax2.set_xlabel('MCP 1 APD')
+        #toRaw = lambda x: x/F
+        #ax2.set_xlim((toRaw(ax.get_xlim()[0]),toRaw(ax.get_xlim()[1])))
+        ax.set_yscale('log')
+        
+        ax = plt.subplot(236)
+        if not use_apd:
+            pulseStart = intstart
+            pulseStop = intstop
+        else:
+            pulseStart = data.attrs['run'].get_array(
+                'SCS_UTC1_ADQ/ADC/1', 'board1.apd.channel_0.pulseStart.value')[0].values
+            pulseStop = data.attrs['run'].get_array(
+                'SCS_UTC1_ADQ/ADC/1', 'board1.apd.channel_0.pulseStop.value')[0].values
+            
+        if 'MCP{}raw'.format(mcp) not in data:
+            tid, data = data.attrs['run'].train_from_index(0)
+            trace = data['SCS_UTC1_ADQ/ADC/1:network']['digitizers.channel_1_D.raw.samples']
+            print('no raw data for MCP{}. Loading trace from MCP1'.format(mcp))
+            label_trace='MCP1 Voltage [V]'
+        else:
+            trace = data['MCP{}raw'.format(mcp)][0]
+            label_trace='MCP{} Voltage [V]'.format(mcp)
+        ax.plot(trace[:pulseStop+25], 'o-', ms=2, label='trace')
+        ax.axvspan(pulseStart, pulseStop, color='C2', alpha=0.2, label='APD region')
+        ax.axvline(pulseStart, color='gray', ls='--')
+        ax.axvline(pulseStop, color='gray', ls='--')
+        ax.set_xlim(pulseStart - 25, pulseStop + 25)
+        ax.set_ylabel(label_trace, size=10)
+        ax.set_xlabel('sample #', size=10)
+        ax.legend(fontsize=8)
+        plt.tight_layout()
+
+    return F