Better finite array retrieval

src/toolbox_scs/base/knife_edge.py

@@ -49,17 +49,13 @@ def prepare_arrays(positions: np.ndarray, intensities: np.ndarray,
     assert positions.shape == intensities.shape
 
     # Clean both arrays by only getting finite values
-    positions, intensities = finite_array(positions, wrt=intensities)
-    intensities, positions = finite_array(intensities, wrt=positions)
+    finite_idx = np.isfinite(positions) & np.isfinite(intensities)
+    positions = positions[finite_idx]
+    intensities = intensities[finite_idx]
 
     return positions, intensities
 
 
-def finite_array(array, *, wrt):
-    index = np.isfinite(wrt)
-    return array[index], wrt[index]
-
-
 def range_mask(array, minimum=None, maximum=None):
     default = np.ones(array.shape, dtype=np.bool)
     min_slice, max_slice = default, default