diff --git a/knife_edge.py b/knife_edge.py index 30a396950b7e7acea0908b963b97852e467ad48f..83d0d8dc83e3ac76c41096a5ccdddc978fcd2103 100644 --- a/knife_edge.py +++ b/knife_edge.py @@ -9,20 +9,26 @@ 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, 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. +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. + 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: - ndarray with beam radius at 1/e^2 in mm and standard error from the fit - in mm. + 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) @@ -42,14 +48,16 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, plo intensities = nrun[signalKey].values.flatten()[sortIdx] if intensities[0] > intensities[-1]: 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 +/- %.1f) mm'%(popt[0], pcov[0,0]**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: @@ -63,6 +71,9 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, plo for lh in leg.legendHandles: lh.set_alpha(1) plt.ylabel(signalKey) - plt.xlabel(axisKey + '-position [mm]') + plt.xlabel(axisKey + ' position [mm]') plt.tight_layout() - return np.array([popt[1], pcov[1,1]**0.5]) \ No newline at end of file + if full: + return popt, pcov + else: + return np.array([popt[1], pcov[1,1]**0.5]) \ No newline at end of file