diff --git a/__init__.py b/__init__.py index 23e93fcfd21dab40e2ef44c354b699046e57c4bb..654b00fd8ddd17321acbc7a68a6d4afe987560bb 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,4 @@ from ToolBox.Load import * from ToolBox.xgm import * from ToolBox.XAS import * +from ToolBox.knife_edge import * diff --git a/knife_edge.py b/knife_edge.py new file mode 100644 index 0000000000000000000000000000000000000000..10659ce8e4881e9869f73ae8ce160d2d3a449b31 --- /dev/null +++ b/knife_edge.py @@ -0,0 +1,82 @@ +""" 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])