Base knife-edge scan analysis implementation

src/toolbox_scs/base/knife_edge.py

+import numpy as np
+from scipy.special import erfc
+from scipy.optimize import curve_fit
+
+
+def knife_edge(positions, intensities, axisRange=None, p0=None):
+    popt, pcov = knife_edge_base(positions, intensities,
+                                 axisRange=axisRange, p0=p0)
+    width, std = popt[1], pcov[1, 1]**0.5
+    return width, std
+
+
+def knife_edge_base(positions, intensities, axisRange=None, p0=None):
+    """
+
+    """
+    # Prepare arrays
+    positions, intensities = prepare_arrays(positions, intensities, axisRange)
+
+    # Estimate a linear slope fitting the data to determine the erfc to use
+    slope = covariance(positions, intensities) / np.var(positions)
+    func = generate_erfc(sign=np.sign(slope) * -1)
+    if p0 is None:
+        p0 = [np.mean(positions), 0.1, np.max(intensities) / 2, 0]
+
+    # Fit
+    try:
+        popt, pcov = curve_fit(func, positions, intensities, p0=p0)
+    except (TypeError, RuntimeError):
+        print("Fit did not converge.")
+        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_slice(positions, *axisRange)
+        positions = positions[slice_]
+        intensities = intensities[slice_]
+
+    # Convert both arrays to 1D of the same size
+    n_pulses = intensities.shape[1]
+    positions = np.repeat(positions, n_pulses)
+    intensities = intensities.flatten()
+    assert positions.shape == intensities.shape
+
+    # Clean both arrays by only getting finite values and sorting
+    positions, intensities = finite_arrays(positions, wrt=intensities)
+    intensities, positions = sort_arrays(intensities, wrt=positions)
+
+    return positions, intensities
+
+
+def sort_arrays(array, *, wrt=None):
+    index = np.argsort(wrt)
+    return array[index], wrt[index]
+
+
+def finite_arrays(array, *, wrt=None):
+    index = np.isfinite(wrt)
+    return array[index], wrt[index]
+
+
+def range_slice(array, minimum=None, maximum=None):
+    default = np.ones(array.shape, dtype=np.bool)
+    min_slice, max_slice = default, default
+    if minimum is not None:
+        if minimum > array[-1]:
+            raise ValueError('The range minimum is too large.')
+        min_slice = array >= minimum
+
+    if maximum is not None:
+        if maximum < array[0]:
+            raise ValueError('The range maximum is too small.')
+        max_slice = array <= maximum
+
+    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 generate_erfc(sign):
+    def func(x, x0, w0, a, b):
+        return a * erfc(sign * np.sqrt(2) * (x - x0) / w0) + b
+
+    return func

src/toolbox_scs/base/tests/test_knife_edge.py

+import numpy as np
+import pytest
+
+from ..knife_edge import range_slice
+
+
+def test_range_slice():
+    arr = np.array([1, 2, 3, 4, 5])
+
+    # Exact
+    slice_ = range_slice(arr, minimum=2, maximum=4)
+    np.testing.assert_array_equal(slice_, [False, True, True, True, False])
+
+    # Range exceeds the closest values
+    slice_ = range_slice(arr, minimum=1.75, maximum=4.25)
+    np.testing.assert_array_equal(slice_, [False, True, True, True, False])
+
+    # Range misses the closest values
+    slice_ = range_slice(arr, minimum=2.25, maximum=3.75)
+    np.testing.assert_array_equal(slice_, [False, False, True, False, False])
+
+    # Equidistant
+    slice_ = range_slice(arr, minimum=2.5, maximum=4.5)
+    np.testing.assert_array_equal(slice_, [False, False, True, True, False])
+
+    # Out of bounds, valid minimum
+    slice_ = range_slice(arr, minimum=0)
+    np.testing.assert_array_equal(slice_, [True, True, True, True, True])
+
+    # Out of bounds, invalid minimum
+    with pytest.raises(ValueError):
+        range_slice(arr, minimum=6)
+
+    # Out of bounds, valid maximum
+    slice_ = range_slice(arr, maximum=6)
+    np.testing.assert_array_equal(slice_, [True, True, True, True, True])
+
+    # Out of bounds, invalid minimum
+    with pytest.raises(ValueError):
+        range_slice(arr, maximum=0)
+
+    # with NaNs
+    arr = np.array([1, np.nan, 3, np.nan, 5])
+    slice_ = range_slice(arr, minimum=3)
+    np.testing.assert_array_equal(slice_, [False, False, True, False, True])