From 20c42486b82d00bc6dfef799d8fda99f72a75f55 Mon Sep 17 00:00:00 2001 From: Laurent Mercadier <laurent.mercadier@xfel.eu> Date: Mon, 18 Nov 2019 21:43:55 +0100 Subject: [PATCH] Adds axisRange parameter, returns fitting function when full=True --- knife_edge.py | 50 +++++++++++++++++++++++++++++++++++--------------- 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/knife_edge.py b/knife_edge.py index 224b8a8..e7b906b 100644 --- a/knife_edge.py +++ b/knife_edge.py @@ -8,10 +8,12 @@ import matplotlib.pyplot as plt import numpy as np from scipy.special import erfc from scipy.optimize import curve_fit +import bisect -def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, full=False, plot=False): +def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', + axisRange=[None,None], 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 + erfc function: f(a,b,u) = a*erfc(u)+b or f(a,b,u) = a*erfc(-u)+b 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 @@ -19,22 +21,24 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, ful 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 + axisRange: list of length 2, minimum and maximum values between which to apply + the fit. + p0: list, initial parameters used for the fit: x0, w0, a, b. 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. + curve_fit as well as the fitting function. 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 integPowerUp(x, x0, w0, a, b): + return a*erfc(-np.sqrt(2)*(x-x0)/w0) + b - def integPowerDown(x, x0, w0, a): - return a*erfc(np.sqrt(2)*(x-x0)/w0) + def integPowerDown(x, x0, w0, a, b): + return a*erfc(np.sqrt(2)*(x-x0)/w0) + b #get the number of pulses per train from the signal source: dim = nrun[signalKey].dims[1] @@ -46,22 +50,38 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, ful sortIdx = np.argsort(positions) positions = positions[sortIdx] intensities = nrun[signalKey].values.flatten()[sortIdx] - + + if axisRange[0] is None or axisRange[0] < positions[0]: + idxMin = 0 + else: + if axisRange[0] >= positions[-1]: + raise ValueError('The minimum value of axisRange is too large') + idxMin = bisect.bisect(positions, axisRange[0]) + if axisRange[1] is None or axisRange[1] > positions[-1]: + idxMax = None + else: + if axisRange[1] <= positions[0]: + raise ValueError('The maximum value of axisRange is too small') + idxMax = bisect.bisect(positions, axisRange[1]) + 1 + positions = positions[idxMin:idxMax] + intensities = intensities[idxMin:idxMax] + # 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)' + funcStr = 'a*erfc(np.sqrt(2)*(x-x0)/w0) + b' else: func = integPowerUp - funcStr = 'a*erfc(-np.sqrt(2)*(x-x0)/w0)' + funcStr = 'a*erfc(-np.sqrt(2)*(x-x0)/w0) + b' if p0 is None: - p0 = [np.mean(positions), 0.1, np.max(intensities)/2] + p0 = [np.mean(positions), 0.1, np.max(intensities)/2, 0] 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)) + print('x0 = (%.3f +/- %.3f) mm'%(popt[0], pcov[0,0]**0.5)) + print('a = %e +/- %e '%(popt[2], pcov[2,2]**0.5)) + print('b = %e +/- %e '%(popt[3], pcov[3,3]**0.5)) if plot: xfit = np.linspace(positions.min(), positions.max(), 1000) @@ -78,6 +98,6 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, ful plt.title(nrun.attrs['runFolder']) plt.tight_layout() if full: - return popt, pcov + return popt, pcov, func else: return np.array([popt[1], pcov[1,1]**0.5]) -- GitLab