Fix time_delta cleaning (#1656)

* Vectorize cleaning using numpy broadcasting

* Fix cleaning behaviour and add unit test

- Ignore unselected neighbor pixels in the time_delta cleaning
- Dont change the original image mask
- Add test case for neighbors arriving inside the time limit, that are
not selected from previous steps

* Fix spelling

- We use the american version everywhere else in the code

* Clean up numpy indexing and sum over boolean array

ctapipe/image/cleaning.py

@@ -232,15 +232,13 @@ def apply_time_delta_cleaning(
     `image[~mask] = 0`
 
     """
-    pixels_to_remove = []
-    mask = mask.copy()  # Create copy so orginal is unchanged
-    for pixel in np.where(mask)[0]:
-        neighbors = geom.neighbor_matrix_sparse[pixel].indices
-        time_diff = np.abs(arrival_times[neighbors] - arrival_times[pixel])
-        if sum(time_diff < time_limit) < min_number_neighbors:
-            pixels_to_remove.append(pixel)
-    mask[pixels_to_remove] = False
-    return mask
+    pixels_to_keep = mask.copy()
+    time_diffs = np.abs(arrival_times[mask, None] - arrival_times)
+    # neighboring pixels arriving in the time limit and previously selected
+    valid_neighbors = (time_diffs < time_limit) & geom.neighbor_matrix[mask] & mask
+    enough_neighbors = np.count_nonzero(valid_neighbors, axis=1) >= min_number_neighbors
+    pixels_to_keep[pixels_to_keep] &= enough_neighbors
+    return pixels_to_keep
 
 
 def fact_image_cleaning(

ctapipe/image/tests/test_cleaning.py

@@ -282,7 +282,7 @@ def test_fact_image_cleaning():
         time_limit=5,
     )
 
-    expected_pixels = np.array([0, 1, 2, 3, 4, 8, 9, 10, 11])
+    expected_pixels = np.array([0, 1, 2, 3, 4, 8, 9])
     expected_mask = np.zeros(len(geom)).astype(bool)
     expected_mask[expected_pixels] = 1
     assert_allclose(mask, expected_mask)
@@ -293,8 +293,8 @@ def test_apply_time_delta_cleaning():
     peak_time = np.zeros(geom.n_pixels, dtype=np.float64)
 
     pixel = 40
-    neighbours = geom.neighbors[pixel]
-    peak_time[neighbours] = 32.0
+    neighbors = geom.neighbors[pixel]
+    peak_time[neighbors] = 32.0
     peak_time[pixel] = 30.0
     mask = peak_time > 0
 
@@ -306,19 +306,29 @@ def test_apply_time_delta_cleaning():
     assert (test_mask == td_mask).all()
 
     # Test time_limit
-    noise_neighbour = neighbours[0]
-    peak_time[noise_neighbour] += 10
+    noise_neighbor = neighbors[0]
+    peak_time[noise_neighbor] += 10
     td_mask = cleaning.apply_time_delta_cleaning(
         geom, mask, peak_time, min_number_neighbors=1, time_limit=5
     )
     test_mask = mask.copy()
-    test_mask[noise_neighbour] = 0
+    test_mask[noise_neighbor] = 0
     assert (test_mask == td_mask).all()
 
-    # Test min_number_neighbours
+    # Test min_number_neighbors
     td_mask = cleaning.apply_time_delta_cleaning(
         geom, mask, peak_time, min_number_neighbors=4, time_limit=5
     )
     test_mask = mask.copy()
-    test_mask[neighbours] = 0
+    test_mask[neighbors] = 0
+    assert (test_mask == td_mask).all()
+
+    # Test unselected neighbors
+    mask[156] = 0
+    peak_time[noise_neighbor] -= 10
+    td_mask = cleaning.apply_time_delta_cleaning(
+        geom, mask, peak_time, min_number_neighbors=3, time_limit=5
+    )
+    test_mask = mask.copy()
+    test_mask[[41, 157]] = 0
     assert (test_mask == td_mask).all()