Use knife edge base in routine function

src/toolbox_scs/base/knife_edge.py

@@ -107,13 +107,7 @@ def prepare_arrays(arrX: np.ndarray, arrY: np.ndarray,
     3. Retrieve finite values.
     """
     # Convert both arrays to 1D of the same size
-    assert arrX.shape[0] == arrY.shape[0]
-    arrX = arrX.flatten()
-    arrY = arrY.flatten()
-    if len(arrX) > len(arrY):
-        arrY = np.repeat(arrY, len(arrX) // len(arrY))
-    else:
-        arrX = np.repeat(arrX, len(arrY) // len(arrX))
+    arrX, arrY = arrays_to1d(arrX, arrY)
 
     # Select ranges
     if xRange is not None:
@@ -133,6 +127,18 @@ def prepare_arrays(arrX: np.ndarray, arrY: np.ndarray,
     return arrX, arrY
 
 
+def arrays_to1d(arrX: np.ndarray, arrY: np.ndarray):
+    """Flatten two arrays and matches their sizes
+    """
+    assert arrX.shape[0] == arrY.shape[0]
+    arrX, arrY = arrX.flatten(), arrY.flatten()
+    if len(arrX) > len(arrY):
+        arrY = np.repeat(arrY, len(arrX) // len(arrY))
+    else:
+        arrX = np.repeat(arrX, len(arrY) // len(arrX))
+    return arrX, arrY
+
+
 def range_mask(array, minimum=None, maximum=None):
     """Retrieve the resulting array from the given minimum and maximum
     """

src/toolbox_scs/routines/knife_edge.py

 """ Toolbox for SCS.
 
     Various utilities function to quickly process data measured
-    at the SCS instruments.
+    at the SCS instrument.
 
     Copyright (2019-) SCS Team.
 """
 import matplotlib.pyplot as plt
 import numpy as np
-from scipy.special import erfc
-from scipy.optimize import curve_fit
-import bisect
+from toolbox_scs.base.knife_edge import knife_edge_base, erfc, arrays_to1d
 
 __all__ = [
     'knife_edge'
 ]
 
 
-def knife_edge(ds, axisKey='scannerX',
-               signalKey='FastADC4peaks',
-               axisRange=[None, None], p0=None,
-               full=False, plot=False):
+def knife_edge(ds, axisKey='scannerX', signalKey='FastADC4peaks',
+               axisRange=[None, None], p0=None, full=False, plot=False,
+               display=False):
     """
     Calculates the beam radius at 1/e^2 from a knife-edge scan by
-    fitting with 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.
+    fitting with erfc function:
+    f(x, x0, w0, a, b) = a*erfc(np.sqrt(2)*(x-x0)/w0) + b
+    with w0 the beam radius at 1/e^2 and x0 the beam center.
 
     Parameters
     ----------
@@ -38,103 +35,62 @@ def knife_edge(ds, axisKey='scannerX',
         edges of the scanning axis between which to apply the fit.
     p0: list of floats, numpy 1D array
         initial parameters used for the fit: x0, w0, a, b. If None, a beam
-        radius of 100 um is assumed.
+        radius of 100 micrometers 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 scipy.optimize.curve_fit as well as the fitting function.
+        matrix from scipy.optimize.curve_fit.
     plot: bool
-        If True, plots the data and the result of the fit.
+        If True, plots the data and the result of the fit. Default is False.
+    display: bool
+        If True, displays info on the fit. True when plot is True, default is
+        False.
 
     Returns
     -------
-    If full is False, ndarray with beam radius at 1/e^2 in mm and standard
+    If full is False, tuple with beam radius at 1/e^2 in mm and standard
         error from the fit in mm. If full is True, returns parameters and
         covariance matrix from scipy.optimize.curve_fit function.
     """
-    def stepUp(x, x0, w0, a, b):
-        return a*erfc(-np.sqrt(2)*(x-x0)/w0) + b
-
-    def stepDown(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 = [k for k in ds[signalKey].dims if k != 'trainId'][0]
-    # duplicate motor position values to match signal shape
-    # this is faster than using ds.stack()
-    positions = np.repeat(ds[axisKey].values,
-                          len(ds[dim])).astype(ds[signalKey].dtype)
-    # sort the data to decide which fitting function to use
-    sortIdx = np.argsort(positions)
-    positions = positions[sortIdx]
-    intensities = ds[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
-    pos_sel = positions[idxMin:idxMax]
-    int_sel = intensities[idxMin:idxMax]
-    no_nan = ~np.isnan(int_sel)
-    pos_sel = pos_sel[no_nan]
-    int_sel = int_sel[no_nan]
+    popt, pcov = knife_edge_base(ds[axisKey].values, ds[signalKey].values,
+                                 axisRange=axisRange, p0=p0)
+    if plot:
+        positions, intensities = arrays_to1d(ds[axisKey].values,
+                                             ds[signalKey].values)
+        plot_knife_edge(positions, intensities, popt, pcov[1, 1]**0.5,
+                        ds.attrs['runFolder'], axisKey, signalKey)
+        display = True
 
-    # estimate a linear slope fitting the data to determine which function
-    # to fit
-    slope = np.cov(pos_sel, int_sel)[0][1]/np.var(pos_sel)
-    if slope < 0:
-        func = stepDown
+    if display:
         funcStr = 'a*erfc(np.sqrt(2)*(x-x0)/w0) + b'
-    else:
-        func = stepUp
-        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]
-    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('w0 = (%.1f +/- %.1f) um' % (np.abs(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 Exception as e:
-        print(f'Could not fit the data with erfc function: {e}.' +
-              ' Try adjusting the axisRange and the initial parameters p0')
-        fitSuccess = False
 
-    if plot:
-        plt.figure(figsize=(7, 4))
-        plt.scatter(positions, intensities, color='C1',
-                    label='exp', s=2, alpha=0.1)
-        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)
-        plt.ylabel(signalKey)
-        plt.xlabel(axisKey + ' position [mm]')
-        plt.title(ds.attrs['runFolder'])
-        plt.tight_layout()
     if full:
-        if fitSuccess:
-            return popt, pcov, func
-        else:
-            return np.zeros(4), np.zeros(2), None
+        return popt, pcov
     else:
-        if fitSuccess:
-            return np.array([popt[1], pcov[1, 1]**0.5])
-        else:
-            return np.zeros(2)
+        return np.abs(popt[1]), pcov[1, 1]**0.5
+
+
+def plot_knife_edge(positions, intensities, fit_params, rel_err, title,
+                    axisKey, signalKey):
+
+    plt.figure(figsize=(7, 4))
+    plt.scatter(positions, intensities, color='C1',
+                label='measured', s=2, alpha=0.1)
+    xfit = np.linspace(positions.min(), positions.max(), 1000)
+    yfit = erfc(xfit, *fit_params)
+    plt.plot(xfit, yfit, color='C4',
+             label=r'fit $\rightarrow$ $w_0=$(%.1f $\pm$ %.1f) $\mu$m' % (
+                                    np.abs(fit_params[1])*1e3, rel_err*1e3))
+    leg = plt.legend()
+    for lh in leg.legendHandles:
+        lh.set_alpha(1)
+    plt.ylabel(signalKey)
+    plt.xlabel(axisKey + ' position [mm]')
+    plt.title(title)
+    plt.tight_layout()