Merge branch 'knife_edge' into 'master'

Knife edge

See merge request SCS/ToolBox!36

__init__.py

 from ToolBox.Load import *
 from ToolBox.xgm import *
 from ToolBox.XAS import *
+from ToolBox.knife_edge import *

knife_edge.py

+""" 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
+from scipy.special import erfc
+from scipy.optimize import curve_fit
+
+def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, full=False, plot=False):
+    ''' Calculates the beam radius at 1/e^2 from a knife-edge scan by fitting with
+        erfc function: f(a, u) = a*erfc(u) or f(a, u) = a*erfc(-u) where 
+        u = sqrt(2)*(x-x0)/w0 with w0 the beam radius at 1/e^2 and x0 the beam center.
+        Inputs:
+            nrun: xarray Dataset containing the detector signal and the motor 
+                  position.
+            axisKey: string, key of the axis against which the knife-edge is 
+                  performed.
+            signalKey: string, key of the detector signal.
+            p0: list, initial parameters used for the fit: x0, w0, a. If None, a beam
+                radius of 100 um is assumed.
+            full: bool: If False, returns the beam radius and standard error. If True,
+                returns the popt, pcov list of parameters and covariance matrix from 
+                curve_fit.
+            plot: bool: If True, plots the data and the result of the fit.
+        Outputs:
+            If full is False, ndarray with beam radius at 1/e^2 in mm and standard 
+            error from the fit in mm. If full is True, returns popt and pcov from 
+            curve_fit function.
+    '''
+    def integPowerUp(x, x0, w0, a):
+        return a*erfc(-np.sqrt(2)*(x-x0)/w0)
+
+    def integPowerDown(x, x0, w0, a):
+        return a*erfc(np.sqrt(2)*(x-x0)/w0)
+
+    #get the number of pulses per train from the signal source:
+    dim = nrun[signalKey].dims[1]
+    #duplicate motor position values to match signal shape
+    #this is much faster than using nrun.stack()
+    positions = np.repeat(nrun[axisKey].values, 
+                          len(nrun[dim])).astype(nrun[signalKey].dtype)
+    #sort the data to decide which fitting function to use
+    sortIdx = np.argsort(positions)
+    positions = positions[sortIdx]
+    intensities = nrun[signalKey].values.flatten()[sortIdx]
+
+    # estimate a linear slope fitting the data to determine which function to fit
+    slope = np.cov(positions, intensities)[0][1]/np.var(positions) 
+    if slope < 0:
+        func = integPowerDown
+        funcStr = 'a*erfc(np.sqrt(2)*(x-x0)/w0)'
+    else:
+        func = integPowerUp
+        funcStr = 'a*erfc(-np.sqrt(2)*(x-x0)/w0)'
+    if p0 is None:
+        p0 = [np.mean(positions), 0.1, np.max(intensities)/2]
+    popt, pcov = curve_fit(func, positions, intensities, p0=p0)
+    print('fitting function:', funcStr)
+    print('w0 = (%.1f +/- %.1f) um'%(popt[1]*1e3, pcov[1,1]**0.5*1e3))
+    print('x0 = (%.3f +/- %.3f) mm'%(popt[0], pcov[0,0]**0.5*1e3))
+    print('a = %e +/- %e '%(popt[2], pcov[2,2]**0.5*1e3))
+    
+    if plot:
+        xfit = np.linspace(positions.min(), positions.max(), 1000)
+        yfit = func(xfit, *popt)
+        plt.figure(figsize=(7,4))
+        plt.scatter(positions, intensities, color='C1', label='exp', s=2, alpha=0.01)
+        plt.plot(xfit, yfit, color='C4', 
+                 label=r'fit $\rightarrow$ $w_0=$(%.1f $\pm$ %.1f) $\mu$m'%(popt[1]*1e3, pcov[1,1]**0.5*1e3))
+        leg = plt.legend()
+        for lh in leg.legendHandles: 
+            lh.set_alpha(1)
+        plt.ylabel(signalKey)
+        plt.xlabel(axisKey + ' position [mm]')
+        plt.tight_layout()
+    if full:
+        return popt, pcov
+    else:
+        return np.array([popt[1], pcov[1,1]**0.5])