Handles case where fit does not work

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)