Support unknown signal dims

src/toolbox_scs/base/knife_edge.py

@@ -43,9 +43,8 @@ def prepare_arrays(positions: np.ndarray, intensities: np.ndarray,
         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()
+    positions = np.repeat(positions, len(intensities) // len(positions))
     assert positions.shape == intensities.shape
 
     # Clean both arrays by only getting finite values

src/toolbox_scs/base/tests/test_knife_edge.py

@@ -45,7 +45,7 @@ def test_range_mask():
     np.testing.assert_array_equal(slice_, [False, False, True, False, True])
 
 
-def test_prepare_arrays():
+def test_prepare_arrays_nans():
     # Setup test values
     trains, pulses = 5, 10
     size = trains * pulses
@@ -74,6 +74,23 @@ def test_prepare_arrays():
     assert np.isfinite(intensities).all()
 
 
+def test_prepare_arrays_size():
+    trains, pulses = 5, 10
+    size = trains * pulses
+    motor = np.arange(trains)
+    signal = np.random.randint(100, size=(trains, pulses))
+
+    # Test finite motor and 2D signals
+    positions, intensities = prepare_arrays(motor, signal)
+    assert positions.shape == (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,)
+
+
 def with_values(array, value, num=5):
     copy = array.astype(np.float)
     copy.ravel()[np.random.choice(copy.size, num, replace=False)] = value