diff --git a/src/toolbox_scs/base/knife_edge.py b/src/toolbox_scs/base/knife_edge.py index e46cf6c2a0eaf85be849e1bef6dd9d81ec78507e..b9dfb696b0ec71861e5ddc91d584eaba2c053a0c 100644 --- a/src/toolbox_scs/base/knife_edge.py +++ b/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 diff --git a/src/toolbox_scs/base/tests/test_knife_edge.py b/src/toolbox_scs/base/tests/test_knife_edge.py index 888afbe4e6f0ee063a5531e6e9dc3f3ddf601076..08457da82917a7386da9da94082f34838feb3e9d 100644 --- a/src/toolbox_scs/base/tests/test_knife_edge.py +++ 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)