More description of fitting parameters

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