diff --git a/knife_edge.py b/knife_edge.py index 099f2f73154b2b5f9db7693445333e33a4ffa5a3..58e0b2bbc9e07670603d28bfb1d7bd70659703d0 100644 --- a/knife_edge.py +++ b/knife_edge.py @@ -76,20 +76,28 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', funcStr = 'a*erfc(-np.sqrt(2)*(x-x0)/w0) + b' if p0 is None: p0 = [np.mean(pos_sel), 0.1, np.max(int_sel)/2, 0] - popt, pcov = curve_fit(func, pos_sel, int_sel, 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)) - print('a = %e +/- %e '%(popt[2], pcov[2,2]**0.5)) - print('b = %e +/- %e '%(popt[3], pcov[3,3]**0.5)) - + try: + popt, pcov = curve_fit(func, pos_sel, int_sel, 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)) + print('a = %e +/- %e '%(popt[2], pcov[2,2]**0.5)) + print('b = %e +/- %e '%(popt[3], pcov[3,3]**0.5)) + fitSuccess = True + except: + print('Could not fit the data with ercf function.' + + ' Try adjusting the axisRange and the initial parameters p0') + fitSuccess = False + 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.1) - 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)) + if fitSuccess: + xfit = np.linspace(positions.min(), positions.max(), 1000) + yfit = func(xfit, *popt) + 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) @@ -98,6 +106,12 @@ def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', plt.title(nrun.attrs['runFolder']) plt.tight_layout() if full: - return popt, pcov, func + if fitSuccess: + return popt, pcov, func + else: + return np.zeros(4), np.zeros(2), None else: - return np.array([popt[1], pcov[1,1]**0.5]) + if fitSuccess: + return np.array([popt[1], pcov[1,1]**0.5]) + else: + return np.zeros(2)