Rename methods, fix and test finite_array

src/toolbox_scs/base/knife_edge.py

@@ -6,7 +6,9 @@ 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
+    width, std = 0, 0
+    if popt is not None and pcov is not None:
+        width, std = popt[1], pcov[1, 1]**0.5
     return width, std
 
 
@@ -36,7 +38,7 @@ def prepare_arrays(positions: np.ndarray, intensities: np.ndarray,
                    axisRange=None):
     # Slice
     if axisRange is not None:
-        slice_ = range_slice(positions, *axisRange)
+        slice_ = range_mask(positions, *axisRange)
         positions = positions[slice_]
         intensities = intensities[slice_]
 
@@ -46,33 +48,28 @@ def prepare_arrays(positions: np.ndarray, intensities: np.ndarray,
     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)
+    # Clean both arrays by only getting finite values
+    positions, intensities = finite_array(positions, wrt=intensities)
+    intensities, positions = finite_array(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):
+def finite_array(array, *, wrt):
     index = np.isfinite(wrt)
     return array[index], wrt[index]
 
 
-def range_slice(array, minimum=None, maximum=None):
+def range_mask(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]:
+        if minimum > np.nanmax(array):
             raise ValueError('The range minimum is too large.')
         min_slice = array >= minimum
 
     if maximum is not None:
-        if maximum < array[0]:
+        if maximum < np.nanmin(array):
             raise ValueError('The range maximum is too small.')
         max_slice = array <= maximum
 

src/toolbox_scs/base/tests/test_knife_edge.py

 import numpy as np
 import pytest
 
-from ..knife_edge import range_slice
+from ..knife_edge import prepare_arrays, range_mask
 
 
-def test_range_slice():
+def test_range_mask():
     arr = np.array([1, 2, 3, 4, 5])
 
     # Exact
-    slice_ = range_slice(arr, minimum=2, maximum=4)
+    slice_ = range_mask(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)
+    slice_ = range_mask(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)
+    slice_ = range_mask(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)
+    slice_ = range_mask(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)
+    slice_ = range_mask(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)
+        range_mask(arr, minimum=6)
 
     # Out of bounds, valid maximum
-    slice_ = range_slice(arr, maximum=6)
+    slice_ = range_mask(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)
+        range_mask(arr, maximum=0)
 
     # with NaNs
     arr = np.array([1, np.nan, 3, np.nan, 5])
-    slice_ = range_slice(arr, minimum=3)
+    slice_ = range_mask(arr, minimum=3)
     np.testing.assert_array_equal(slice_, [False, False, True, False, True])
+
+
+def test_prepare_arrays():
+    # Setup test values
+    trains, pulses = 5, 10
+    size = trains * pulses
+    motor = np.arange(trains)
+    signal = np.random.randint(100, size=(trains, pulses))
+
+    # Test finite motor and signal values
+    positions, intensities = prepare_arrays(motor, signal)
+    assert positions.shape == (size,)
+    assert intensities.shape == (size,)
+
+    # Test finite motors and signals with NaNs
+    signal_nan = with_values(signal, value=np.nan, num=20)
+    positions, intensities = prepare_arrays(motor, signal_nan)
+    assert positions.shape == (size-20,)
+    assert np.isfinite(positions).all()
+    assert intensities.shape == (size-20,)
+    assert np.isfinite(intensities).all()
+
+    # Test finite signals and motors with NaNs
+    motor_nan = with_values(motor, value=np.nan, num=3)
+    positions, intensities = prepare_arrays(motor_nan, signal)
+    assert positions.shape == ((trains-3) * pulses,)
+    assert np.isfinite(positions).all()
+    assert intensities.shape == ((trains-3) * pulses,)
+    assert np.isfinite(intensities).all()
+
+
+def with_values(array, value, num=5):
+    copy = array.astype(np.float)
+    copy.ravel()[np.random.choice(copy.size, num, replace=False)] = value
+    return copy