Add generic fit function, generalize prepare_arrays()

src/toolbox_scs/base/knife_edge.py

@@ -8,7 +8,7 @@ def knife_edge(positions, intensities, axisRange=None, p0=None):
                                  axisRange=axisRange, p0=p0)
     width, std = 0, 0
     if popt is not None and pcov is not None:
-        width, std = popt[1], pcov[1, 1]**0.5
+        width, std = np.abs(popt[1]), pcov[1, 1]**0.5
     return width, std
 
 
@@ -17,40 +17,58 @@ def knife_edge_base(positions, intensities, axisRange=None, p0=None):
 
     """
     # Prepare arrays
-    positions, intensities = prepare_arrays(positions, intensities, axisRange)
+    positions, intensities = prepare_arrays(positions, intensities,
+                                            xRange=axisRange)
 
     # Estimate initial fitting params
     if p0 is None:
         p0 = [np.mean(positions), 0.1, np.max(intensities) / 2, 0]
 
+    # Fit
+    popt, pcov = function_fit(erfc, positions, intensities, p0=p0)
+
+    return popt, pcov
+
+
+def function_fit(func, x, y, **kwargs):
+    """A wrapper for scipy.optimize curve_fit()
+    """
     # Fit
     try:
-        popt, pcov = curve_fit(erfc, positions, intensities, p0=p0)
-    except (TypeError, RuntimeError):
-        print("Fit did not converge.")
+        popt, pcov = curve_fit(func, x, y, **kwargs)
+    except (TypeError, RuntimeError) as err:
+        print("Fit did not converge:", err)
         popt, pcov = (None, None)
     return popt, pcov
 
 
-def prepare_arrays(positions: np.ndarray, intensities: np.ndarray,
-                   axisRange=None):
-    # Slice
-    if axisRange is not None:
-        slice_ = range_mask(positions, *axisRange)
-        positions = positions[slice_]
-        intensities = intensities[slice_]
-
+def prepare_arrays(arrX: np.ndarray, arrY: np.ndarray,
+                   xRange=None, yRange=None):
     # Convert both arrays to 1D of the same size
-    intensities = intensities.flatten()
-    positions = np.repeat(positions, len(intensities) // len(positions))
-    assert positions.shape == intensities.shape
+    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))
+
+    # Select ranges
+    if xRange is not None:
+        mask_ = range_mask(arrX, *xRange)
+        arrX = arrX[mask_]
+        arrY = arrY[mask_]
+    if yRange is not None:
+        mask_ = range_mask(arrY, *yRange)
+        arrX = arrX[mask_]
+        arrY = arrY[mask_]
 
     # Clean both arrays by only getting finite values
-    finite_idx = np.isfinite(positions) & np.isfinite(intensities)
-    positions = positions[finite_idx]
-    intensities = intensities[finite_idx]
+    finite_idx = np.isfinite(arrX) & np.isfinite(arrY)
+    arrX = arrX[finite_idx]
+    arrY = arrY[finite_idx]
 
-    return positions, intensities
+    return arrX, arrY
 
 
 def range_mask(array, minimum=None, maximum=None):
@@ -69,10 +87,5 @@ def range_mask(array, minimum=None, maximum=None):
     return min_slice & max_slice
 
 
-def covariance(a: np.ndarray, b: np.ndarray):
-    assert a.ndim == b.ndim == 1
-    return np.cov(a, b)[0][1]
-
-
 def erfc(x, x0, w0, a, b):
     return a * special.erfc(np.sqrt(2) * (x - x0) / w0) + b

src/toolbox_scs/base/tests/test_knife_edge.py

@@ -86,9 +86,9 @@ def test_prepare_arrays_size():
     assert intensities.shape == (size,)
 
     # Test finite motor and 1D signals
-    positions, intensities = prepare_arrays(motor, signal.reshape(1, -1))
-    assert positions.shape == (size,)
-    assert intensities.shape == (size,)
+    positions, intensities = prepare_arrays(motor, signal[:, 0])
+    assert positions.shape == (trains,)
+    assert intensities.shape == (trains,)
 
 
 def test_knife_edge_base():
@@ -96,14 +96,14 @@ def test_knife_edge_base():
     x = np.linspace(-3, 3)
     y = erfc(x, *p0)
     noise = y * np.random.normal(0, .02, y.shape)  # 2% error
-    eff_y = (y + noise).reshape(1, -1)
+    eff_y = y + noise
 
     # Test noisy data
     popt, _ = knife_edge_base(x, eff_y)
     np.testing.assert_allclose(p0, popt, atol=1e-1)
 
     # Test flipped data
-    popt, _ = knife_edge_base(x, eff_y[:, ::-1])
+    popt, _ = knife_edge_base(x, eff_y[::-1])
     p0[1] = abs(p0[1])  # Absolute value when flipped
     np.testing.assert_allclose(p0, popt, atol=1e-1)